Compare commits
118 Commits
cc603aba06
...
feat/proje
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cc0e258e8c | ||
|
|
12e203e63d | ||
|
|
ffcd7390f0 | ||
|
|
91e880f9d4 | ||
|
|
7d47ca54be | ||
|
|
956fa88853 | ||
|
|
fb2f59ccea | ||
|
|
56dbb7f4cd | ||
|
|
506f517851 | ||
|
|
520c186991 | ||
|
|
582bf27deb | ||
|
|
2aeb453229 | ||
|
|
b7a4edac90 | ||
|
|
822b4cd8b1 | ||
|
|
ab24fc4c91 | ||
|
|
a98e99f7a2 | ||
|
|
a0ff285bcd | ||
|
|
177c1a87dd | ||
|
|
441a4ea05c | ||
|
|
a693a64bf5 | ||
|
|
67562b8092 | ||
|
|
6f4c68b359 | ||
|
|
c20c6d7853 | ||
|
|
6787e690ba | ||
|
|
cb8f56d909 | ||
|
|
2c7cac9e03 | ||
|
|
ea9094f47f | ||
|
|
d5fea95561 | ||
|
|
0b5ef48463 | ||
|
|
ca8721e1ac | ||
|
|
f658e5e6a3 | ||
|
|
341ee140e5 | ||
|
|
741b9b87fb | ||
|
|
2d8abb6311 | ||
|
|
e668e3fd20 | ||
|
|
7ccdad431f | ||
|
|
4073863dc6 | ||
|
|
a85f8fde29 | ||
|
|
90500a3462 | ||
|
|
c1a8ac7669 | ||
|
|
c510cbaae5 | ||
|
|
ce139bbac3 | ||
|
|
3cf067faea | ||
|
|
7253f6fe72 | ||
|
|
41db3a7089 | ||
|
|
cc94194fd1 | ||
|
|
96c91e386d | ||
|
|
c0aef71141 | ||
|
|
467abc8d42 | ||
|
|
5753f8def9 | ||
|
|
e672b58b6f | ||
|
|
d8add7e8cb | ||
|
|
c6c4578f9a | ||
|
|
3aa0b36a6c | ||
|
|
fa231a3642 | ||
|
|
d91c98f86d | ||
|
|
c0619f5c4d | ||
|
|
da282229ff | ||
|
|
7fa6ad5760 | ||
|
|
dcd14220ca | ||
|
|
3cc32569d9 | ||
|
|
bf445ac2ce | ||
|
|
a2d6d689e4 | ||
|
|
aa8bcbf0d8 | ||
|
|
1ce1d492b0 | ||
|
|
552b8eb305 | ||
|
|
0d93b3960d | ||
|
|
f07580574b | ||
|
|
1a8bf11f90 | ||
|
|
e7cdce8287 | ||
|
|
58bc6efd4b | ||
|
|
6c450805cb | ||
|
|
f340d0fa3e | ||
|
|
edc53cb6eb | ||
|
|
725cece5c1 | ||
|
|
297e20ce8d | ||
|
|
5a03bd1cfb | ||
|
|
87b7a1c6c9 | ||
|
|
826f64d6bb | ||
| 5faa6b1d7c | |||
| 02a9684cd6 | |||
| fae9efee0d | |||
| 30b062dd4a | |||
| 2a0331d7ce | |||
| 13fd8677c1 | |||
| 9bd629cb59 | |||
| 9c97702daa | |||
| a1e364c9c0 | |||
| 5b55f1292a | |||
| 5bc9ea6cd6 | |||
| f7404b6f66 | |||
| d667e43c73 | |||
| fe085a7951 | |||
| 2de67213f8 | |||
| f6ed383b3a | |||
| 9332e29e53 | |||
| 618076193a | |||
| 34f01234c9 | |||
| 0bd46937d3 | |||
| e6b5bc2e7d | |||
| c90ed58078 | |||
| 76c8f2bdad | |||
| 393b3befd6 | |||
| 2c08275934 | |||
| 7cb384fa63 | |||
| 7efaeba283 | |||
| b61ded8458 | |||
| ac71d99f9a | |||
| 3b3b3baf25 | |||
| 45415bb9ee | |||
| a775a2da18 | |||
| 24772f2b67 | |||
| fd1396a710 | |||
| 914f70bd85 | |||
| 608d6c784f | |||
| 19ad5be97f | |||
| 1dfd088e18 | |||
| c6e1e4e7fd |
85
.env.example
85
.env.example
@@ -2,7 +2,7 @@
|
|||||||
ENV=dev
|
ENV=dev
|
||||||
|
|
||||||
# ── Database ──────────────────────────────────────────────────────────────────
|
# ── Database ──────────────────────────────────────────────────────────────────
|
||||||
DATABASE_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/adiuva
|
DATABASE_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/adiuvai
|
||||||
|
|
||||||
# ── Auth ──────────────────────────────────────────────────────────────────────
|
# ── Auth ──────────────────────────────────────────────────────────────────────
|
||||||
JWT_SECRET=replace-with-a-long-random-secret
|
JWT_SECRET=replace-with-a-long-random-secret
|
||||||
@@ -13,31 +13,82 @@ JWT_REFRESH_TOKEN_EXPIRE_DAYS=30
|
|||||||
# ── LLM ───────────────────────────────────────────────────────────────────────
|
# ── LLM ───────────────────────────────────────────────────────────────────────
|
||||||
# LiteLLM model identifiers — change to swap providers without code changes.
|
# LiteLLM model identifiers — change to swap providers without code changes.
|
||||||
# Examples: gpt-4o, anthropic/claude-sonnet-4-20250514, gemini/gemini-pro, ollama/llama3
|
# Examples: gpt-4o, anthropic/claude-sonnet-4-20250514, gemini/gemini-pro, ollama/llama3
|
||||||
|
#
|
||||||
|
# API keys — only the key(s) matching your chosen provider(s) are required.
|
||||||
|
# The correct key is picked automatically from the model prefix (e.g.
|
||||||
|
# "anthropic/..." → ANTHROPIC_API_KEY, "gemini/..." → GOOGLE_API_KEY).
|
||||||
OPENAI_API_KEY=
|
OPENAI_API_KEY=
|
||||||
ANTHROPIC_API_KEY=
|
ANTHROPIC_API_KEY=
|
||||||
GOOGLE_API_KEY=
|
GOOGLE_API_KEY=
|
||||||
LLM_MODEL=gpt-4o
|
CEREBRAS_API_KEY=
|
||||||
LLM_ROUTER_MODEL=gpt-4o-mini
|
GROQ_API_KEY=
|
||||||
|
DEEPSEEK_API_KEY=
|
||||||
|
|
||||||
|
# Default model used by any agent that does not have a specific override below.
|
||||||
|
LLM_MODEL=gpt-5-mini
|
||||||
|
LLM_EMBED_MODEL=text-embedding-3-small
|
||||||
|
|
||||||
|
# GitHub Copilot — leave empty to use the LiteLLM default token directory.
|
||||||
|
# In Docker, point this to a named-volume path so tokens survive restarts.
|
||||||
|
# GITHUB_COPILOT_TOKEN_DIR=
|
||||||
|
|
||||||
|
# ── Per-agent model overrides ─────────────────────────────────────────────────
|
||||||
|
# Leave a value empty to fall back to LLM_MODEL.
|
||||||
|
# Each agent resolves its API key from the model prefix automatically.
|
||||||
|
#
|
||||||
|
# Intent classifier — routes user messages to the right domain agent.
|
||||||
|
# A small/fast model (e.g. gpt-4o-mini) is usually sufficient here.
|
||||||
|
LLM_MODEL_CLASSIFIER=
|
||||||
|
|
||||||
|
# Home-agent — handles chat from the home screen (all tools available).
|
||||||
|
LLM_MODEL_HOME_AGENT=
|
||||||
|
|
||||||
|
# Floating-agent — handles contextual chat triggered from a task/project/note.
|
||||||
|
LLM_MODEL_FLOATING_AGENT=
|
||||||
|
|
||||||
|
# Unified-processor — processes local directory files (local agent runner).
|
||||||
|
LLM_MODEL_UNIFIED_PROCESSOR=
|
||||||
|
|
||||||
|
# Cloud-processor — fetches and processes data from cloud connectors.
|
||||||
|
LLM_MODEL_CLOUD_PROCESSOR=
|
||||||
|
|
||||||
|
# Brief-agent — produces home and project text briefs.
|
||||||
|
# A small model (e.g. gpt-4o-mini) is sufficient.
|
||||||
|
# LLM_MODEL_BRIEF_AGENT=
|
||||||
|
|
||||||
|
# Task-brief-agent — per-task deep research (Stage 1 executive assistant).
|
||||||
|
# Needs tool-use + reasoning; a capable model recommended (e.g. gpt-4o, gemini-2.5-flash).
|
||||||
|
# LLM_MODEL_TASK_BRIEF_AGENT=
|
||||||
|
|
||||||
|
# Setup-agent — guided journey to build an AgentConfig via WebSocket chat.
|
||||||
|
LLM_MODEL_SETUP_AGENT=
|
||||||
|
|
||||||
|
# Memory-extractor — Mem0-style extract/decide pipeline (Phase 2).
|
||||||
|
# Defaults to gpt-4o-mini when empty (fast + cheap, temperature=0).
|
||||||
|
LLM_MODEL_MEMORY_EXTRACTOR=
|
||||||
|
|
||||||
|
# Memory-miner — proactive pattern mining from episodic history (Phase 5, Power+ only).
|
||||||
|
# Defaults to gpt-4o-mini when empty.
|
||||||
|
LLM_MODEL_MEMORY_MINER=
|
||||||
|
|
||||||
|
# Memory-auditor — weekly contradiction scan + relation label canonicalization (Phase 7).
|
||||||
|
# Defaults to LLM_MODEL when empty (a reasoning-capable model is recommended).
|
||||||
|
LLM_MODEL_MEMORY_AUDITOR=
|
||||||
|
|
||||||
|
# Scheduler — set to false to disable memory cron jobs (automatically false in tests).
|
||||||
|
SCHEDULER_ENABLED=true
|
||||||
|
|
||||||
# ── Stripe (leave empty to stub billing) ──────────────────────────────────────
|
# ── Stripe (leave empty to stub billing) ──────────────────────────────────────
|
||||||
STRIPE_SECRET_KEY=
|
STRIPE_SECRET_KEY=
|
||||||
STRIPE_WEBHOOK_SECRET=
|
STRIPE_WEBHOOK_SECRET=
|
||||||
|
|
||||||
# ── AWS / S3 ──────────────────────────────────────────────────────────────────
|
|
||||||
S3_BUCKET=adiuva
|
|
||||||
S3_REGION=us-east-1
|
|
||||||
S3_ENDPOINT_URL=
|
|
||||||
AWS_ACCESS_KEY_ID=
|
|
||||||
AWS_SECRET_ACCESS_KEY=
|
|
||||||
# For MinIO (homelab): S3_ENDPOINT_URL=http://minio:9000
|
|
||||||
|
|
||||||
# ── Vector Store ──────────────────────────────────────────────────────────────
|
# ── Langfuse (leave empty to disable observability) ───────────────────────────
|
||||||
# Pinecone is used when PINECONE_API_KEY is set; otherwise falls back to Qdrant.
|
LANGFUSE_SECRET_KEY=
|
||||||
PINECONE_API_KEY=
|
LANGFUSE_PUBLIC_KEY=
|
||||||
PINECONE_INDEX=adiuva
|
# LANGFUSE_BASE_URL=https://cloud.langfuse.com # EU (default)
|
||||||
QDRANT_URL=
|
# LANGFUSE_BASE_URL=https://us.cloud.langfuse.com # US
|
||||||
QDRANT_API_KEY=
|
# LANGFUSE_BASE_URL=http://localhost:3000 # Self-hosted
|
||||||
# For local Qdrant (homelab): QDRANT_URL=http://qdrant:6333
|
|
||||||
|
|
||||||
# ── CORS ──────────────────────────────────────────────────────────────────────
|
# ── CORS ──────────────────────────────────────────────────────────────────────
|
||||||
# Comma-separated list parsed by Settings (override default if needed)
|
# Comma-separated list parsed by Settings (override default if needed)
|
||||||
|
|||||||
@@ -48,23 +48,23 @@ jobs:
|
|||||||
key: ${{ secrets.SSH_KEY }}
|
key: ${{ secrets.SSH_KEY }}
|
||||||
script: |
|
script: |
|
||||||
set -e
|
set -e
|
||||||
DEPLOY_DIR="/opt/adiuva-api"
|
DEPLOY_DIR="/opt/adiuvai-api"
|
||||||
REPO_URL="http://10.0.0.119:3000/${{ gitea.repository }}.git"
|
REPO_URL="http://10.0.0.119:3000/${{ gitea.repository }}.git"
|
||||||
TAG="${{ gitea.ref_name }}"
|
TAG="${{ gitea.ref_name }}"
|
||||||
|
|
||||||
# ── Pull latest code ──
|
# ── Pull latest code ──
|
||||||
cd /tmp && rm -rf adiuva-api-deploy
|
cd /tmp && rm -rf adiuvai-api-deploy
|
||||||
git clone --depth 1 --branch "${TAG}" "${REPO_URL}" adiuva-api-deploy
|
git clone --depth 1 --branch "${TAG}" "${REPO_URL}" adiuvai-api-deploy
|
||||||
|
|
||||||
# ── Sync source (preserve .env) ──
|
# ── Sync source (preserve .env) ──
|
||||||
cp -rf /tmp/adiuva-api-deploy/app/ \
|
cp -rf /tmp/adiuvai-api-deploy/app/ \
|
||||||
/tmp/adiuva-api-deploy/alembic/ \
|
/tmp/adiuvai-api-deploy/alembic/ \
|
||||||
/tmp/adiuva-api-deploy/alembic.ini \
|
/tmp/adiuvai-api-deploy/alembic.ini \
|
||||||
/tmp/adiuva-api-deploy/Dockerfile \
|
/tmp/adiuvai-api-deploy/Dockerfile \
|
||||||
/tmp/adiuva-api-deploy/docker-compose.yml \
|
/tmp/adiuvai-api-deploy/docker-compose.yml \
|
||||||
/tmp/adiuva-api-deploy/requirements.txt \
|
/tmp/adiuvai-api-deploy/requirements.txt \
|
||||||
"$DEPLOY_DIR/"
|
"$DEPLOY_DIR/"
|
||||||
rm -rf /tmp/adiuva-api-deploy
|
rm -rf /tmp/adiuvai-api-deploy
|
||||||
|
|
||||||
# ── Verify .env ──
|
# ── Verify .env ──
|
||||||
if [ ! -f "$DEPLOY_DIR/.env" ]; then
|
if [ ! -f "$DEPLOY_DIR/.env" ]; then
|
||||||
|
|||||||
4
.github/workflows/ci.yml
vendored
4
.github/workflows/ci.yml
vendored
@@ -58,7 +58,7 @@ jobs:
|
|||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
|
|
||||||
- name: Build image
|
- name: Build image
|
||||||
run: docker build -t adiuva-api:ci .
|
run: docker build -t adiuvai-api:ci .
|
||||||
|
|
||||||
- name: Verify gunicorn installed
|
- name: Verify gunicorn installed
|
||||||
run: docker run --rm adiuva-api:ci gunicorn --version
|
run: docker run --rm adiuvai-api:ci gunicorn --version
|
||||||
|
|||||||
5
.gitignore
vendored
5
.gitignore
vendored
@@ -21,13 +21,18 @@ env/
|
|||||||
.pytest_cache/
|
.pytest_cache/
|
||||||
htmlcov/
|
htmlcov/
|
||||||
.coverage
|
.coverage
|
||||||
|
tests/fixtures/private*/
|
||||||
|
|
||||||
# Docker
|
# Docker
|
||||||
*.log
|
*.log
|
||||||
|
|
||||||
# OS
|
# OS
|
||||||
.DS_Store
|
.DS_Store
|
||||||
|
|
||||||
|
# Smoke scripts (dev-only, not for CI)
|
||||||
|
scripts/smoke_*.py
|
||||||
Thumbs.db
|
Thumbs.db
|
||||||
|
|
||||||
# Claude Code
|
# Claude Code
|
||||||
.claude/
|
.claude/
|
||||||
|
logs/
|
||||||
|
|||||||
@@ -1,243 +0,0 @@
|
|||||||
# AI Refactor Plan — Adiuva Backend
|
|
||||||
|
|
||||||
> **Objective:** Transform backend tools from JSON-action-descriptor-returning functions into real bidirectional executors. Each tool sends structured CRUD operations to the Electron client via WebSocket, receives real data back, and returns meaningful results to the LLM. The LLM reasons about actual user data instead of serialized action payloads.
|
|
||||||
>
|
|
||||||
> **Electron app:** Lives at `../adiuva/`. See `../adiuva/AI_REFACTOR_PLAN.md`.
|
|
||||||
>
|
|
||||||
> **Protocol:** Execute steps sequentially. Each step is atomic and committable. Mark `[x]` when done.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Architecture — Before vs After
|
|
||||||
|
|
||||||
### Before (current)
|
|
||||||
```
|
|
||||||
LLM calls list_tasks(status="todo")
|
|
||||||
→ tool returns: '{"action":"list","table":"tasks","filters":{"status":"todo"}}'
|
|
||||||
→ _tool_loop feeds that JSON string as ToolMessage to LLM
|
|
||||||
→ LLM sees a descriptor, NOT real data — cannot reason about tasks
|
|
||||||
→ Final response: generic "Here are your tasks" (no actual task data)
|
|
||||||
→ Action descriptors sent in final WS frame for Electron to execute post-response
|
|
||||||
```
|
|
||||||
|
|
||||||
### After (target)
|
|
||||||
```
|
|
||||||
LLM calls list_tasks(status="todo")
|
|
||||||
→ tool calls execute_on_client(action="select", table="tasks", filters={status:"todo"})
|
|
||||||
→ WS frame sent to Electron: {type:"tool_call", id:"abc", action:"select", table:"tasks", filters:{status:"todo"}}
|
|
||||||
→ Electron runs: db.select().from(tasks).where(eq(tasks.status, "todo")).all()
|
|
||||||
→ WS frame back: {type:"tool_result", id:"abc", rows:[{id:"1",title:"Buy milk",...}, ...]}
|
|
||||||
→ tool returns: "Found 3 tasks: 1. Buy milk (high, due tomorrow) 2. ..."
|
|
||||||
→ _tool_loop feeds that as ToolMessage to LLM
|
|
||||||
→ LLM sees REAL data — can reason, count, compare, summarize
|
|
||||||
```
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## WS Protocol — Typed Frames
|
|
||||||
|
|
||||||
| Direction | `type` | Payload |
|
|
||||||
|---|---|---|
|
|
||||||
| Client → Server | `chat_request` | `{ message: str, context: ChatContext }` |
|
|
||||||
| Server → Client | `text_chunk` | `{ text: str }` |
|
|
||||||
| Server → Client | `tool_call` | `{ id: str, action: str, table?: str, data?: dict, filters?: dict, vector?: list[float], limit?: int }` |
|
|
||||||
| Client → Server | `tool_result` | `{ id: str, row?: dict, rows?: list[dict], results?: list[dict], deleted?: bool, ok?: bool, error?: str }` |
|
|
||||||
| Server → Client | `final` | `{ response: str }` |
|
|
||||||
| Server → Client | `ping` | `{}` |
|
|
||||||
|
|
||||||
**Actions:**
|
|
||||||
|
|
||||||
| `action` | What Electron does (Drizzle) | `tool_result` shape |
|
|
||||||
|---|---|---|
|
|
||||||
| `select` | `db.select().from(table).where(filters)` | `{ rows: [...] }` |
|
|
||||||
| `get` | `db.select().from(table).where(id=...).get()` | `{ row: {...} or null }` |
|
|
||||||
| `insert` | `db.insert(table).values({id: uuid(), ...data}).returning().get()` | `{ row: {...} }` |
|
|
||||||
| `update` | `db.update(table).set(updates).where(id=...).returning().get()` | `{ row: {...} }` |
|
|
||||||
| `delete` | `db.delete(table).where(id=...).run()` | `{ deleted: true }` |
|
|
||||||
| `vector_upsert` | LanceDB upsert with pre-computed vector | `{ ok: true }` |
|
|
||||||
| `vector_search` | LanceDB search by vector | `{ results: [{id, content, score}...] }` |
|
|
||||||
|
|
||||||
**Electron generates IDs + timestamps.** Backend tools never send `id` or `createdAt` in `insert` data — Electron adds `id: uuid()`, `createdAt: Date.now()`, `updatedAt: Date.now()`.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## SQLite Schema Reference (Electron's local database)
|
|
||||||
|
|
||||||
Tools must use **camelCase** field names (Drizzle maps them to snake_case internally):
|
|
||||||
|
|
||||||
| Table | Columns |
|
|
||||||
|---|---|
|
|
||||||
| `tasks` | id, projectId, title, description, status (todo\|in_progress\|done), priority (high\|medium\|low), assignee (JSON array string), dueDate (ms), isAiSuggested (0\|1), isApproved (0\|1), createdAt (ms) |
|
|
||||||
| `projects` | id, clientId, name, status (active\|archived), aiSummary, createdAt (ms) |
|
|
||||||
| `checkpoints` | id, projectId (required), title, date (ms), isAiSuggested (0\|1), isApproved (0\|1), createdAt (ms) |
|
|
||||||
| `notes` | id, projectId, title, content (markdown), createdAt (ms), updatedAt (ms) |
|
|
||||||
| `taskComments` | id, taskId, author, content, createdAt (ms) |
|
|
||||||
| `clients` | id, parentId, name, industry, createdAt (ms) |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Phase B — Backend Changes
|
|
||||||
|
|
||||||
### Step B.1 — WS context + frame types
|
|
||||||
- [x] Create `app/core/ws_context.py` (~25 lines):
|
|
||||||
- `_client_executor: ContextVar[Callable]` — holds the async callback for the current WS session
|
|
||||||
- `async def execute_on_client(action, table=None, data=None, filters=None, vector=None, limit=None) -> dict`:
|
|
||||||
- Reads callback from ContextVar
|
|
||||||
- Builds `tool_call` payload: `{id: str(uuid4()), action, table, data, filters, vector, limit}` (omits None fields)
|
|
||||||
- Calls `await callback(payload)` — which sends the WS frame and waits for `tool_result`
|
|
||||||
- Returns the result dict
|
|
||||||
- `def set_client_executor(fn)` / `def clear_client_executor()` — ContextVar management
|
|
||||||
- [x] Add to `app/schemas.py`:
|
|
||||||
- `WsFrameType(str, Enum)`: `chat_request`, `text_chunk`, `tool_call`, `tool_result`, `final`, `ping`
|
|
||||||
- `WsToolCall(BaseModel)`: `type`, `id`, `action`, `table?`, `data?`, `filters?`, `vector?`, `limit?`
|
|
||||||
- `WsToolResult(BaseModel)`: `type`, `id`, `row?`, `rows?`, `results?`, `deleted?`, `ok?`, `error?`
|
|
||||||
- `WsTextChunk(BaseModel)`: `type`, `text`
|
|
||||||
- `WsFinal(BaseModel)`: `type`, `response`
|
|
||||||
- **Files:** `app/core/ws_context.py`, `app/schemas.py`
|
|
||||||
- **Outcome:** Any tool can `await execute_on_client(...)` to query/mutate the user's local DB.
|
|
||||||
|
|
||||||
### Step B.2 — Rewrite all 23 tools to use `execute_on_client()`
|
|
||||||
- [x] Each tool: same `@tool` decorator, same parameters, same docstring. Replace `return json.dumps({...})` body with:
|
|
||||||
1. Call `result = await execute_on_client(action=..., table=..., data/filters=...)`
|
|
||||||
2. Return human-readable string with confirmation + key data from `result`
|
|
||||||
|
|
||||||
- [x] **`app/agents/task_agent.py` (8 tools):**
|
|
||||||
- `list_tasks(project_id, status, search, order_by)`:
|
|
||||||
```python
|
|
||||||
result = await execute_on_client(action="select", table="tasks", filters={
|
|
||||||
"projectId": project_id or None,
|
|
||||||
"status": status or None,
|
|
||||||
"search": search or None,
|
|
||||||
"orderBy": order_by or None,
|
|
||||||
})
|
|
||||||
rows = result.get("rows", [])
|
|
||||||
if not rows:
|
|
||||||
return "No tasks found matching the given filters."
|
|
||||||
lines = [f"- {r['title']} (status: {r['status']}, priority: {r['priority']}, id: {r['id']})" for r in rows]
|
|
||||||
return f"Found {len(rows)} task(s):\n" + "\n".join(lines)
|
|
||||||
```
|
|
||||||
- `create_task(title, ...)`:
|
|
||||||
```python
|
|
||||||
result = await execute_on_client(action="insert", table="tasks", data={
|
|
||||||
"title": title, "description": description or None, "status": status,
|
|
||||||
"priority": priority, "assignee": assignees, "dueDate": due_date or None,
|
|
||||||
"projectId": project_id or None, "isAiSuggested": is_ai_suggested, "isApproved": is_approved,
|
|
||||||
})
|
|
||||||
row = result["row"]
|
|
||||||
return f"Task created: '{row['title']}' (id: {row['id']}, status: {row['status']}, priority: {row['priority']})"
|
|
||||||
```
|
|
||||||
- `update_task(task_id, ...)`: build updates dict (same logic as now) → `execute_on_client(action="update", table="tasks", data={"id": task_id, "updates": updates})` → return "Task updated: {title}"
|
|
||||||
- `delete_task(task_id)`: `execute_on_client(action="delete", table="tasks", data={"id": task_id})` → return "Task deleted"
|
|
||||||
- `list_tasks_due_today()`: calculate today's start/end ms → `execute_on_client(action="select", table="tasks", filters={"dueDateFrom": start, "dueDateTo": end})` → format + return
|
|
||||||
- `list_task_comments(task_id)`: `execute_on_client(action="select", table="taskComments", filters={"taskId": task_id})` → format + return
|
|
||||||
- `add_task_comment(task_id, author, content)`: `execute_on_client(action="insert", table="taskComments", data={...})` → return confirmation
|
|
||||||
- `delete_task_comment(comment_id)`: `execute_on_client(action="delete", table="taskComments", data={"id": comment_id})` → return confirmation
|
|
||||||
|
|
||||||
- [x] **`app/agents/project_agent.py` (6 tools):**
|
|
||||||
- `list_projects(client_id, include_archived)`: `execute_on_client(action="select", table="projects", filters={clientId, includeArchived})` → format + return
|
|
||||||
- `list_all_projects()`: `execute_on_client(action="select", table="projects")` → format + return
|
|
||||||
- `get_project(project_id)`: `execute_on_client(action="get", table="projects", data={"id": project_id})` → return project details or "not found"
|
|
||||||
- `create_project(name, client_id)`: `execute_on_client(action="insert", table="projects", data={name, clientId})` → return confirmation + id
|
|
||||||
- `update_project(project_id, ...)`: build updates → `execute_on_client(action="update", ...)` → return confirmation
|
|
||||||
- `delete_project(project_id)`: `execute_on_client(action="delete", ...)` → return confirmation
|
|
||||||
|
|
||||||
- [x] **`app/agents/checkpoint_agent.py` (4 tools):**
|
|
||||||
- `list_checkpoints(project_id)`: `execute_on_client(action="select", table="checkpoints", filters={projectId})` → format + return
|
|
||||||
- `create_checkpoint(project_id, title, date, ...)`: `execute_on_client(action="insert", table="checkpoints", data={...})` → return confirmation + id
|
|
||||||
- `update_checkpoint(checkpoint_id, ...)`: build updates → `execute_on_client(action="update", ...)` → return confirmation
|
|
||||||
- `delete_checkpoint(checkpoint_id)`: `execute_on_client(action="delete", ...)` → return confirmation
|
|
||||||
|
|
||||||
- [x] **`app/agents/note_agent.py` (5 tools):**
|
|
||||||
- `list_notes(project_id)`: `execute_on_client(action="select", table="notes", filters={projectId})` → format + return
|
|
||||||
- `get_note(note_id)`: `execute_on_client(action="get", table="notes", data={"id": note_id})` → return full content or "not found"
|
|
||||||
- `create_note(title, content, project_id)`: `execute_on_client(action="insert", table="notes", data={...})` → then `execute_on_client(action="vector_upsert", data={id, projectId, content}, vector=await embed(content))` → return confirmation
|
|
||||||
- `update_note(note_id, ...)`: build updates → `execute_on_client(action="update", ...)` → then vector_upsert for updated content → return confirmation
|
|
||||||
- `delete_note(note_id)`: `execute_on_client(action="delete", ...)` → return confirmation
|
|
||||||
|
|
||||||
- **Files:** `app/agents/task_agent.py`, `app/agents/project_agent.py`, `app/agents/checkpoint_agent.py`, `app/agents/note_agent.py`
|
|
||||||
- **Outcome:** All 23 tools query real user data via WS. LLM sees actual rows, not action descriptors.
|
|
||||||
|
|
||||||
### Step B.3 — Bidirectional WebSocket handler
|
|
||||||
- [x] Refactor `app/api/routes/chat.py` WS endpoint:
|
|
||||||
- After auth + accept + receive `chat_request`:
|
|
||||||
1. Create `execute_on_client` callback closure capturing the websocket:
|
|
||||||
```python
|
|
||||||
pending_calls: dict[str, asyncio.Future] = {}
|
|
||||||
|
|
||||||
async def on_client_result(frame: dict):
|
|
||||||
"""Called when a tool_result frame arrives from Electron."""
|
|
||||||
fut = pending_calls.pop(frame["id"], None)
|
|
||||||
if fut and not fut.done():
|
|
||||||
fut.set_result(frame)
|
|
||||||
|
|
||||||
async def execute_callback(payload: dict) -> dict:
|
|
||||||
"""Send tool_call to Electron, wait for tool_result."""
|
|
||||||
call_id = payload["id"]
|
|
||||||
fut = asyncio.get_event_loop().create_future()
|
|
||||||
pending_calls[call_id] = fut
|
|
||||||
await websocket.send_text(json.dumps({"type": "tool_call", **payload}))
|
|
||||||
return await asyncio.wait_for(fut, timeout=30.0)
|
|
||||||
```
|
|
||||||
2. Set `client_executor` ContextVar with `execute_callback`
|
|
||||||
3. Run orchestrator in a task — it calls agents, agents call tools, tools call `execute_on_client()` which goes through the callback
|
|
||||||
4. In parallel, run a message receive loop that dispatches incoming frames:
|
|
||||||
- `tool_result` → `on_client_result(frame)`
|
|
||||||
- `ping` → ignore
|
|
||||||
5. Orchestrator yields `text_chunk` frames → send to client
|
|
||||||
6. Send `final` frame when done
|
|
||||||
7. Clear ContextVar
|
|
||||||
- Keep heartbeat ping every 30s
|
|
||||||
- 30s timeout on `tool_result` — if Electron doesn't respond, future raises `TimeoutError`, tool returns error string to LLM
|
|
||||||
- **Files:** `app/api/routes/chat.py`
|
|
||||||
- **Outcome:** Full bidirectional WS. Tool calls and text streaming happen concurrently on the same connection.
|
|
||||||
|
|
||||||
### Step B.4 — `_tool_loop` — no changes needed
|
|
||||||
- [x] Verify `app/core/agent_registry.py` works unchanged:
|
|
||||||
- `_tool_loop` calls `tool_fn.ainvoke(args)` → tool awaits `execute_on_client()` (WS round-trip) → returns string → `ToolMessage(content=string)` → LLM sees real data
|
|
||||||
- The async WS round-trip happens inside each tool. `_tool_loop` just sees an awaited tool returning a string — same as before, different content.
|
|
||||||
- **No code changes.** Just verify + add a log line for tool execution times if desired.
|
|
||||||
|
|
||||||
### Step B.5 — Orchestrator cleanup
|
|
||||||
- [x] Update `app/core/orchestrator.py`:
|
|
||||||
- `orchestrate_stream()`: remove `"actions": []` from final frame. Final becomes: `{"done": true, "response": "..."}`
|
|
||||||
- No other changes — `classify_intent` → `call_agent` → chunk response → final frame
|
|
||||||
- **Files:** `app/core/orchestrator.py`
|
|
||||||
- **Outcome:** Clean final frame. No more action descriptors in the protocol.
|
|
||||||
|
|
||||||
### Step B.6 — Add `/vectors/embed` endpoint
|
|
||||||
- [x] Add to `app/api/routes/vectors.py`:
|
|
||||||
- `POST /api/v1/storage/vectors/embed`:
|
|
||||||
- Request: `{ text: str }`
|
|
||||||
- Response: `{ vector: list[float] }` (1536-dim from `text-embedding-3-small`)
|
|
||||||
- Auth required (JWT)
|
|
||||||
- Used by:
|
|
||||||
- Backend tools: `note_agent` calls this before `vector_upsert`
|
|
||||||
- Electron: `vectordb.ts` calls this for note embedding on create/update
|
|
||||||
- **Files:** `app/api/routes/vectors.py`
|
|
||||||
- **Outcome:** Single embedding endpoint. Both backend tools and Electron can generate vectors.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Verification
|
|
||||||
|
|
||||||
| What to test | How |
|
|
||||||
|---|---|
|
|
||||||
| **Read flow** | "List my tasks" → `list_tasks` → `tool_call{select, tasks}` → Electron returns rows → LLM describes real tasks |
|
|
||||||
| **Write flow** | "Create a task called Buy milk" → `create_task` → `tool_call{insert, tasks, data:{title:"Buy milk"}}` → Electron inserts + returns row → tool confirms with id |
|
|
||||||
| **Multi-tool** | "How many todo tasks do I have?" → `list_tasks(status=todo)` → LLM counts actual rows → "You have 3 todo tasks" |
|
|
||||||
| **Vector search** | "Find notes about deployment" → tool embeds → `tool_call{vector_search, vector:[...]}` → Electron searches LanceDB → returns matching notes |
|
|
||||||
| **Vector upsert** | "Create a note about..." → insert note → vector_upsert with embedding → both SQLite + LanceDB updated |
|
|
||||||
| **Tool timeout** | Disconnect Electron mid-conversation → 30s timeout → tool returns error → LLM handles gracefully |
|
|
||||||
| **Concurrent calls** | Agent calls 2 tools in sequence → each does WS round-trip → both succeed → LLM sees both results |
|
|
||||||
| **_tool_loop max iter** | Verify 5-iteration limit still works → after 5 tool calls, LLM forced to answer without tools |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Execution Notes
|
|
||||||
|
|
||||||
- **Phase 1 is the critical path.** Auth + backend client + drizzle executor + orchestrator refactor must land first.
|
|
||||||
- **Steps 1.1–1.4 are additive** — existing app keeps working until Step 1.5 swaps the orchestrator.
|
|
||||||
- **Step 2.1 is the point of no return** — after removing LangChain, there's no local AI fallback.
|
|
||||||
- **Phase B (backend changes) must land before Phase 1.3–1.5** — Electron needs the bidirectional WS to talk to.
|
|
||||||
- **Phase 3 and Phase 4 are independent** — can be parallelized after Phase 2.
|
|
||||||
- **One step at a time.** Mark `[x]` and commit with `step N.N complete: <outcome>`.
|
|
||||||
533
BACKEND_PLAN.md
533
BACKEND_PLAN.md
@@ -1,533 +0,0 @@
|
|||||||
# Backend Plan — Adiuva Cloud API
|
|
||||||
|
|
||||||
> **Separate repository.** This document defines the FastAPI backend that the Electron app communicates with.
|
|
||||||
>
|
|
||||||
> The backend owns: orchestration logic, chat agent intelligence, prompt IP, auth, billing, E2E backup blob storage, cloud storage (encrypted blobs), cloud vector store, and plugin marketplace.
|
|
||||||
> The backend NEVER persists user data in plaintext. Cloud storage blobs are E2E encrypted before upload — the backend only verifies integrity, never decrypts.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Project Structure
|
|
||||||
|
|
||||||
```
|
|
||||||
adiuva-api/
|
|
||||||
├── app/
|
|
||||||
│ ├── __init__.py
|
|
||||||
│ ├── main.py # FastAPI entry + CORS + lifespan + router includes
|
|
||||||
│ ├── core/
|
|
||||||
│ │ ├── __init__.py
|
|
||||||
│ │ ├── agent_registry.py # Base classes + singleton registry
|
|
||||||
│ │ ├── orchestrator.py # LLM-based intent router
|
|
||||||
│ │ ├── execution_plan.py # Plan builder + cache
|
|
||||||
│ │ └── plugin_loader.py # Dynamic agent loading
|
|
||||||
│ ├── agents/ # Chat agents (proprietary logic + prompts)
|
|
||||||
│ │ ├── __init__.py # Auto-registers all agents
|
|
||||||
│ │ ├── task_agent.py
|
|
||||||
│ │ ├── calendar_agent.py
|
|
||||||
│ │ ├── email_agent.py
|
|
||||||
│ │ └── analytics_agent.py
|
|
||||||
│ ├── api/
|
|
||||||
│ │ ├── __init__.py
|
|
||||||
│ │ ├── routes/
|
|
||||||
│ │ │ ├── __init__.py
|
|
||||||
│ │ │ ├── chat.py # POST /chat + WS /chat/stream
|
|
||||||
│ │ │ ├── plans.py # GET /plans/playbook
|
|
||||||
│ │ │ ├── storage.py # CRUD cloud storage (E2E encrypted blobs)
|
|
||||||
│ │ │ ├── vectors.py # Upsert/search cloud vector store
|
|
||||||
│ │ │ ├── backup.py # PUT/GET /backup
|
|
||||||
│ │ │ ├── plugins.py # Plugin marketplace
|
|
||||||
│ │ │ ├── auth.py # Register/login/refresh
|
|
||||||
│ │ │ └── billing.py # Checkout/webhook/subscription
|
|
||||||
│ │ └── middleware/
|
|
||||||
│ │ ├── __init__.py
|
|
||||||
│ │ ├── auth.py # JWT validation
|
|
||||||
│ │ ├── rate_limit.py # Tier-aware rate limiting
|
|
||||||
│ │ └── sanitizer.py # Strip prompt metadata from responses
|
|
||||||
│ ├── storage/
|
|
||||||
│ │ ├── __init__.py
|
|
||||||
│ │ ├── blob_store.py # S3 for E2E encrypted blobs
|
|
||||||
│ │ ├── vector_store.py # Cloud vector store (Pinecone/Qdrant)
|
|
||||||
│ │ └── encryption.py # Integrity verification only — NO decryption
|
|
||||||
│ ├── marketplace/
|
|
||||||
│ │ ├── __init__.py
|
|
||||||
│ │ ├── plugin_registry.py # Plugin catalog (metadata, versions, ratings)
|
|
||||||
│ │ ├── plugin_review.py # Review queue + approval workflow
|
|
||||||
│ │ └── revenue_share.py # 70/30 split tracking with Stripe Connect
|
|
||||||
│ ├── billing/
|
|
||||||
│ │ ├── __init__.py
|
|
||||||
│ │ ├── stripe_service.py # Stripe checkout + webhooks
|
|
||||||
│ │ └── tier_manager.py # Feature matrix per tier
|
|
||||||
│ └── config/
|
|
||||||
│ ├── __init__.py
|
|
||||||
│ └── settings.py # Pydantic BaseSettings (env-based)
|
|
||||||
├── tests/
|
|
||||||
│ ├── __init__.py
|
|
||||||
│ ├── conftest.py # Fixtures: test client, mock agents, mock LLM
|
|
||||||
│ ├── test_orchestrator.py
|
|
||||||
│ ├── test_agents.py
|
|
||||||
│ ├── test_auth.py
|
|
||||||
│ ├── test_backup.py
|
|
||||||
│ ├── test_storage.py
|
|
||||||
│ └── test_plugins.py
|
|
||||||
├── alembic/ # DB migrations (auth/billing/marketplace tables only)
|
|
||||||
│ ├── alembic.ini
|
|
||||||
│ └── versions/
|
|
||||||
├── requirements.txt
|
|
||||||
├── Dockerfile
|
|
||||||
├── docker-compose.yml # App + PostgreSQL + Redis (dev)
|
|
||||||
├── .env.example
|
|
||||||
└── README.md
|
|
||||||
```
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Step-by-Step Implementation
|
|
||||||
|
|
||||||
### Step 1 — Project scaffolding ✅
|
|
||||||
- [x] Initialize repo with the directory structure above
|
|
||||||
- [x] Write `requirements.txt`:
|
|
||||||
```
|
|
||||||
fastapi>=0.115.0
|
|
||||||
uvicorn[standard]>=0.34.0
|
|
||||||
langchain>=0.3.0
|
|
||||||
langchain-openai>=0.3.0
|
|
||||||
pydantic>=2.10.0
|
|
||||||
python-jose[cryptography]>=3.3.0
|
|
||||||
stripe>=11.0.0
|
|
||||||
boto3>=1.35.0
|
|
||||||
slowapi>=0.1.9
|
|
||||||
sqlalchemy>=2.0.0
|
|
||||||
asyncpg>=0.30.0
|
|
||||||
alembic>=1.14.0
|
|
||||||
bcrypt>=4.2.0
|
|
||||||
python-dotenv>=1.0.0
|
|
||||||
httpx>=0.28.0
|
|
||||||
websockets>=14.0
|
|
||||||
pytest>=8.0.0
|
|
||||||
pytest-asyncio>=0.24.0
|
|
||||||
```
|
|
||||||
- [x] Write `app/main.py`: FastAPI app with CORS (allow `app://`, `http://localhost:*`), lifespan (init DB pool, init agent registry), include all routers under `/api/v1`
|
|
||||||
- [x] Write `app/config/settings.py`: `Settings(BaseSettings)` with fields: `DATABASE_URL`, `JWT_SECRET`, `JWT_ALGORITHM` (default HS256), `STRIPE_SECRET_KEY`, `STRIPE_WEBHOOK_SECRET`, `S3_BUCKET`, `S3_REGION`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `OPENAI_API_KEY`, `CORS_ORIGINS`, `ENV` (dev/prod), `PINECONE_API_KEY`, `PINECONE_INDEX`, `QDRANT_URL`, `QDRANT_API_KEY`
|
|
||||||
- [x] Write `Dockerfile`: Python 3.12 slim, multi-stage (builder + runtime), non-root user
|
|
||||||
- [x] Write `docker-compose.yml`: app, postgres:16, optional redis
|
|
||||||
- [x] Write `.env.example`
|
|
||||||
- **Outcome:** Runnable FastAPI skeleton (returns 404 on all routes).
|
|
||||||
|
|
||||||
### Step 2 — Pydantic schemas (API contracts) ✅
|
|
||||||
- [x] Create `app/schemas.py` (mirrors `src/shared/api-types.ts` from Electron repo):
|
|
||||||
- `ChatRequest`: `message: str`, `context: ChatContext`, `execution_mode: Literal['direct', 'plan']`
|
|
||||||
- `ChatContext`: `user_profile: dict`, `relevant_documents: list[str]`, `recent_tasks: list[dict]`, `conversation_history: list[dict]`
|
|
||||||
- `ChatResponse`: `response: str`, `actions: list[PlanAction]`
|
|
||||||
- `PlanAction`: `type: Literal['create_record', 'update_record', 'delete_record', 'index_document', 'send_notification', 'call_agent']`, `table: str | None`, `data: dict | None`, `agent: str | None`
|
|
||||||
- `ExecutionPlan`: `agent: str`, `steps: list[PlanStep]`
|
|
||||||
- `PlanStep`: `action: str`, `prompt_template: str | None`, `variables: dict | None`, `data_from_step: int | None`
|
|
||||||
- `BackupMetadata`: `version: int`, `timestamp: int`, `checksum: str`, `chunk_count: int`
|
|
||||||
- `BillingTier`: `Literal['free', 'pro', 'power', 'team']`
|
|
||||||
- `AuthTokens`: `access_token: str`, `refresh_token: str`, `expires_at: int`
|
|
||||||
- `UserProfile`: `id: str`, `email: str`, `tier: BillingTier`
|
|
||||||
- `StorageRecord`: `id: str`, `user_id: str`, `table: str`, `blob: bytes`, `checksum: str`, `created_at: int`, `updated_at: int` — blob is always E2E encrypted by client
|
|
||||||
- `StorageRecordCreate`: `table: str`, `blob: bytes`, `checksum: str`
|
|
||||||
- `StorageRecordUpdate`: `blob: bytes`, `checksum: str`
|
|
||||||
- `VectorUpsertRequest`: `vectors: list[VectorItem]`
|
|
||||||
- `VectorItem`: `id: str`, `blob: bytes`, `checksum: str` — vector + metadata encrypted by client
|
|
||||||
- `VectorSearchRequest`: `query_blob: bytes`, `top_k: int = 10`
|
|
||||||
- `VectorSearchResponse`: `results: list[VectorSearchResult]`
|
|
||||||
- `VectorSearchResult`: `id: str`, `score: float`, `blob: bytes`
|
|
||||||
- `PluginManifest`: `id: str`, `name: str`, `description: str`, `version: str`, `author: str`, `permissions: list[str]`, `category: str`, `price_cents: int = 0`
|
|
||||||
- `PluginListResponse`: `plugins: list[PluginManifest]`, `total: int`, `page: int`
|
|
||||||
- `PluginInstallRequest`: `plugin_id: str`
|
|
||||||
- **Outcome:** All request/response models defined and validated.
|
|
||||||
|
|
||||||
### Step 3 — Agent Registry + base classes ✅
|
|
||||||
- [x] `app/core/agent_registry.py`:
|
|
||||||
- `BaseAgent(ABC)`:
|
|
||||||
- `user_id: str`, `shared_memory: dict`, `vector_store_context: list[str]`, `skills: list[str]`
|
|
||||||
- Abstract `get_name() -> str`, `get_description() -> str`
|
|
||||||
- `ChatAgent(BaseAgent)`:
|
|
||||||
- Abstract `async handle(query: str, context: dict) -> str`
|
|
||||||
- Abstract `get_tools() -> list` (LangChain tool definitions)
|
|
||||||
- Concrete `_tool_loop(llm, messages, tools, max_iter=5) -> str` — shared tool-calling loop
|
|
||||||
- `AgentRegistry` (singleton):
|
|
||||||
- `_agents: dict[str, ChatAgent]`
|
|
||||||
- `register(agent_class)` — decorator pattern
|
|
||||||
- `get(name) -> ChatAgent`
|
|
||||||
- `list_agents() -> list[dict]` — returns `[{name, description}]` for orchestrator prompt
|
|
||||||
- `async call_agent(name, query, context) -> str` — for inter-agent calls
|
|
||||||
- [x] Unit tests: register, get, list, call_agent with mock
|
|
||||||
- **Outcome:** Pluggable agent framework.
|
|
||||||
|
|
||||||
### Step 4 — Orchestrator ✅
|
|
||||||
- [x] `app/core/orchestrator.py`:
|
|
||||||
- `async classify_intent(message, context, registry) -> str`:
|
|
||||||
- System prompt: "You are an intent classifier. Given the user message and context, decide which agent to route to. Available agents: {registry.list_agents()}. Respond with just the agent name."
|
|
||||||
- Uses gpt-4o-mini via LangChain for low latency
|
|
||||||
- Falls back to `task_agent` if no clear match
|
|
||||||
- `async route_single(agent_name, message, context) -> ChatResponse`:
|
|
||||||
- Instantiates agent from registry
|
|
||||||
- Calls `agent.handle(message, context)`
|
|
||||||
- Returns response + any actions the agent produced
|
|
||||||
- `async route_pipeline(agent_names, message, context) -> ChatResponse`:
|
|
||||||
- Executes agents in sequence
|
|
||||||
- Each agent receives `{...context, previous_results: [...]}`
|
|
||||||
- Final synthesis via LLM: "Summarize these agent results into a coherent response"
|
|
||||||
- `async orchestrate(request: ChatRequest) -> ChatResponse | ExecutionPlan`:
|
|
||||||
- Main entry point
|
|
||||||
- Context is transparent to orchestrator — data may originate from local or cloud storage on the client side
|
|
||||||
- Classifies intent
|
|
||||||
- If `execution_mode == 'direct'`: route + return response
|
|
||||||
- If `execution_mode == 'plan'`: route + return execution plan with template IDs
|
|
||||||
- `async orchestrate_stream(request: ChatRequest) -> AsyncGenerator[str, None]`:
|
|
||||||
- Same as orchestrate but yields tokens for WebSocket streaming
|
|
||||||
- [x] Integration tests with mocked LLM and mocked agents
|
|
||||||
- **Outcome:** Intelligent routing with single-agent and pipeline modes.
|
|
||||||
|
|
||||||
### Step 5 — Execution Plan generator ✅
|
|
||||||
- [x] `app/core/execution_plan.py`:
|
|
||||||
- `PromptTemplateRegistry`: dict of `template_id -> prompt_text`. Templates are server-side only — client receives IDs.
|
|
||||||
- `ExecutionPlanBuilder`:
|
|
||||||
- `add_step(action, params) -> self`
|
|
||||||
- `add_llm_step(template_id, variables) -> self`
|
|
||||||
- `add_data_step(action, data_from_step) -> self`
|
|
||||||
- `build() -> ExecutionPlan` — validates step references
|
|
||||||
- `PlanCache`:
|
|
||||||
- In-memory LRU (maxsize=1000)
|
|
||||||
- `cache_plan(key, plan)`, `get_plan(key)`, `get_all_playbooks() -> list[ExecutionPlan]`
|
|
||||||
- Playbooks are pre-built plans for common operations (e.g., "create task from email", "generate weekly report")
|
|
||||||
- **Outcome:** Plans are cacheable as playbooks. Prompt IP never leaves the server.
|
|
||||||
|
|
||||||
### Step 6 — Chat Agents ✅
|
|
||||||
- [x] `app/agents/task_agent.py` — `@registry.register`:
|
|
||||||
- Description: "Manages tasks and comments: list, create, update, delete, due-today, comments"
|
|
||||||
- Tools (8): `list_tasks(project_id, status, search, order_by)`, `create_task(title, description, status, priority, assignees, due_date, project_id, is_ai_suggested, is_approved)`, `update_task(task_id, ...)`, `delete_task(task_id)`, `list_tasks_due_today()`, `list_task_comments(task_id)`, `add_task_comment(task_id, author, content)`, `delete_task_comment(comment_id)`
|
|
||||||
- status: `todo|in_progress|done`; priority: `high|medium|low`; assignees: JSON-encoded string; due_date: ms timestamp
|
|
||||||
- Accepts flexible context; sentinel `-1` for optional integer update fields
|
|
||||||
- [x] `app/agents/checkpoint_agent.py` — `@registry.register`:
|
|
||||||
- Description: "Manages project checkpoints (milestones): list, create, update, delete"
|
|
||||||
- Tools (4): `list_checkpoints(project_id)`, `create_checkpoint(project_id, title, date, is_ai_suggested, is_approved)`, `update_checkpoint(checkpoint_id, ...)`, `delete_checkpoint(checkpoint_id)`
|
|
||||||
- `project_id` is required for create; date is a ms timestamp; supports AI-suggestion + approval workflow
|
|
||||||
- [x] `app/agents/project_agent.py` — `@registry.register`:
|
|
||||||
- Description: "Manages projects: list, get, create, update, archive, delete"
|
|
||||||
- Tools (6): `list_projects(client_id, include_archived)`, `list_all_projects()`, `get_project(project_id)`, `create_project(name, client_id)`, `update_project(project_id, ...)`, `delete_project(project_id)`
|
|
||||||
- status: `active|archived`; prefers archive over deletion (docstring guard on delete)
|
|
||||||
- [x] `app/agents/note_agent.py` — `@registry.register`:
|
|
||||||
- Description: "Manages notes: list, get, create, update, delete"
|
|
||||||
- Tools (5): `list_notes(project_id)`, `get_note(note_id)`, `create_note(title, content, project_id)`, `update_note(note_id, ...)`, `delete_note(note_id)`
|
|
||||||
- content is Markdown; `get_note` should be called before update to preserve existing content
|
|
||||||
- [x] `app/agents/__init__.py`: imports all four agent modules to trigger `@registry.register` decorators
|
|
||||||
- [x] Unit tests per agent with mocked LLM (registration, names, tool counts, handle(), direct tool invocation)
|
|
||||||
- **Outcome:** Four domain-specific agents matching the UI data model (Tasks, Checkpoints, Projects, Notes), all registered and tested.
|
|
||||||
|
|
||||||
### Step 7 — Storage Layer ✅
|
|
||||||
- [x] `app/storage/blob_store.py`:
|
|
||||||
- `BlobStore`: `async upload`, `async download`, `async delete` (idempotent), `async list_keys`
|
|
||||||
- Keys: `{user_id}/{table}/{record_id}` — backend never inspects blob content
|
|
||||||
- boto3 S3 with SSE-S3 at-rest encryption; client checksum stored in S3 object metadata
|
|
||||||
- [x] `app/storage/vector_store.py`:
|
|
||||||
- `VectorStore`: `async upsert`, `async search`, `async delete`
|
|
||||||
- Pinecone (default, `namespace=user_id`) or Qdrant (`user_id` payload filter) — runtime-configurable
|
|
||||||
- 32-dim SHA-256-derived float vector; blob stored as base64 in metadata/payload
|
|
||||||
- ANN on encrypted data: known accuracy trade-off, documented
|
|
||||||
- [x] `app/storage/encryption.py`:
|
|
||||||
- `verify_checksum(blob, checksum) -> bool` — SHA-256 + `hmac.compare_digest` (constant-time)
|
|
||||||
- `reject_if_tampered(blob, checksum)` — raises `HTTP 400` on mismatch
|
|
||||||
- Backend NEVER holds decryption keys
|
|
||||||
- [x] `app/schemas.py`: added `StorageRecord*`, `VectorItem`, `VectorUpsertRequest`, `VectorSearch*`, `Plugin*` schemas
|
|
||||||
- [x] `app/config/settings.py`: added `PINECONE_API_KEY`, `PINECONE_INDEX`, `QDRANT_URL`, `QDRANT_API_KEY`
|
|
||||||
- [x] `requirements.txt`: added `moto[s3]`, `pinecone`, `qdrant-client`
|
|
||||||
- [x] 37 unit tests covering encryption, BlobStore (moto), VectorStore Pinecone, VectorStore Qdrant
|
|
||||||
- **Outcome:** Cloud storage layer that handles E2E encrypted blobs without ever accessing plaintext.
|
|
||||||
|
|
||||||
### Step 8 — API Routes ✅
|
|
||||||
|
|
||||||
#### 8a — Chat endpoint
|
|
||||||
- [x] `app/api/routes/chat.py`:
|
|
||||||
- `POST /api/v1/chat`:
|
|
||||||
- Request: `ChatRequest`
|
|
||||||
- Calls `orchestrate(request)` or `orchestrate()` + `build_plan()`
|
|
||||||
- Response: `ChatResponse` or `ExecutionPlan`
|
|
||||||
- `WebSocket /api/v1/chat/stream`:
|
|
||||||
- Client sends `ChatRequest` as first JSON frame
|
|
||||||
- Server yields token strings via `orchestrate_stream()`
|
|
||||||
- Final frame: JSON `ChatResponse` with `{"done": true, "response": "...", "actions": [...]}`
|
|
||||||
- Heartbeat ping every 30s to keep connection alive
|
|
||||||
|
|
||||||
#### 8b — Plans endpoint
|
|
||||||
- [x] `app/api/routes/plans.py`:
|
|
||||||
- `GET /api/v1/plans/playbook`: Returns all playbooks available for the user's tier
|
|
||||||
- `GET /api/v1/plans/playbook/{plan_id}`: Returns a specific plan
|
|
||||||
|
|
||||||
#### 8c — Storage endpoint (cloud records)
|
|
||||||
- [x] `app/api/routes/storage.py`:
|
|
||||||
- `POST /api/v1/storage/records`: Create encrypted record
|
|
||||||
- Request: `StorageRecordCreate`
|
|
||||||
- Verifies checksum, stores blob in S3, inserts metadata row in PostgreSQL
|
|
||||||
- Response: `{id: str, created_at: int}`
|
|
||||||
- `GET /api/v1/storage/records`: List record metadata (no blobs)
|
|
||||||
- Query params: `table: str`, `page: int`, `limit: int`
|
|
||||||
- Response: `list[{id, table, checksum, created_at, updated_at}]`
|
|
||||||
- `GET /api/v1/storage/records/{id}`: Download encrypted blob
|
|
||||||
- Response: blob bytes + `X-Checksum` header
|
|
||||||
- `PUT /api/v1/storage/records/{id}`: Update encrypted blob
|
|
||||||
- Request: `StorageRecordUpdate`
|
|
||||||
- `DELETE /api/v1/storage/records/{id}`: Delete record + S3 blob
|
|
||||||
- All routes enforce tier cloud_storage_gb quota via `TierManager.check_quota(user_id)`
|
|
||||||
|
|
||||||
#### 8d — Vectors endpoint (cloud vector store)
|
|
||||||
- [x] `app/api/routes/vectors.py`:
|
|
||||||
- `POST /api/v1/storage/vectors/upsert`:
|
|
||||||
- Request: `VectorUpsertRequest`
|
|
||||||
- Verifies checksums, delegates to `VectorStore.upsert()`
|
|
||||||
- Response: `{upserted: int}`
|
|
||||||
- `POST /api/v1/storage/vectors/search`:
|
|
||||||
- Request: `VectorSearchRequest`
|
|
||||||
- Delegates to `VectorStore.search()`
|
|
||||||
- Response: `VectorSearchResponse`
|
|
||||||
- `DELETE /api/v1/storage/vectors`:
|
|
||||||
- Request: `{ids: list[str]}`
|
|
||||||
|
|
||||||
#### 8e — Backup endpoint
|
|
||||||
- [x] `app/api/routes/backup.py`:
|
|
||||||
- `PUT /api/v1/backup`: Accepts binary blob + metadata headers (`X-Backup-Version`, `X-Backup-Timestamp`, `X-Backup-Checksum`). Stores in S3 keyed by `{user_id}/{timestamp}`. Enforces tier limits:
|
|
||||||
- Free: 0 (no backup)
|
|
||||||
- Pro: 5 GB
|
|
||||||
- Power: 25 GB
|
|
||||||
- Team: unlimited
|
|
||||||
- `GET /api/v1/backup`: Returns latest blob for authenticated user. Supports `If-Modified-Since`.
|
|
||||||
- `GET /api/v1/backup/history`: Returns list of `BackupMetadata` (no blobs).
|
|
||||||
- `DELETE /api/v1/backup/{backup_id}`: Delete specific backup.
|
|
||||||
|
|
||||||
#### 8f — Plugins endpoint
|
|
||||||
- [x] `app/api/routes/plugins.py`:
|
|
||||||
- `GET /api/v1/plugins`:
|
|
||||||
- Query params: `category: str | None`, `q: str | None`, `page: int`, `sort: Literal['rating', 'installs', 'newest']`
|
|
||||||
- Response: `PluginListResponse`
|
|
||||||
- Available from Power tier and above
|
|
||||||
- `GET /api/v1/plugins/{id}`:
|
|
||||||
- Response: `PluginManifest` + ratings + install count
|
|
||||||
- `POST /api/v1/plugins/{id}/install`:
|
|
||||||
- Request: `PluginInstallRequest`
|
|
||||||
- Records installation for the user (billing tracking, analytics)
|
|
||||||
- If plugin is paid: triggers Stripe Connect charge + revenue split (70% developer, 30% platform)
|
|
||||||
- Response: `{ok: true, download_url: str}` — signed S3 URL for plugin package
|
|
||||||
- `DELETE /api/v1/plugins/{id}/install`:
|
|
||||||
- Unregisters installation
|
|
||||||
|
|
||||||
#### 8g — Auth endpoint
|
|
||||||
- [x] `app/api/routes/auth.py`:
|
|
||||||
- `POST /api/v1/auth/register`: `{email, password}` → bcrypt hash → insert user → return `AuthTokens`
|
|
||||||
- `POST /api/v1/auth/login`: Validate credentials → return `AuthTokens`
|
|
||||||
- `POST /api/v1/auth/refresh`: Rotate refresh token → return new `AuthTokens`
|
|
||||||
- `GET /api/v1/auth/me`: Return `UserProfile` for current JWT
|
|
||||||
|
|
||||||
#### 8h — Billing endpoint
|
|
||||||
- [x] `app/api/routes/billing.py`:
|
|
||||||
- `POST /api/v1/billing/checkout`: Creates Stripe checkout session → returns URL
|
|
||||||
- `POST /api/v1/billing/webhook`: Handles Stripe webhooks (subscription lifecycle)
|
|
||||||
- `GET /api/v1/billing/subscription`: Returns current subscription info
|
|
||||||
- `DELETE /api/v1/billing/subscription`: Cancels subscription
|
|
||||||
|
|
||||||
- **Outcome:** Complete REST + WebSocket API covering orchestration, storage, vectors, backup, marketplace.
|
|
||||||
|
|
||||||
### Step 9 — Middleware
|
|
||||||
|
|
||||||
#### 9a — Auth middleware
|
|
||||||
- [x] `app/api/middleware/auth.py`:
|
|
||||||
- FastAPI dependency: `get_current_user(token: str = Depends(oauth2_scheme)) -> UserProfile`
|
|
||||||
- Validates JWT signature, expiry, extracts `user_id` and `tier`
|
|
||||||
- Raises `401` on invalid/expired token
|
|
||||||
- Exempt routes: `/api/v1/auth/register`, `/api/v1/auth/login`, `/api/v1/billing/webhook`
|
|
||||||
|
|
||||||
#### 9b — Rate limiter
|
|
||||||
- [x] `app/api/middleware/rate_limit.py`:
|
|
||||||
- Uses `slowapi` with `Limiter(key_func=get_user_id_from_jwt)`
|
|
||||||
- Tier-based limits:
|
|
||||||
- Free: 20 req/min
|
|
||||||
- Pro: 60 req/min
|
|
||||||
- Power: 120 req/min
|
|
||||||
- Team: 200 req/seat/min
|
|
||||||
- Custom 429 response with `Retry-After` header
|
|
||||||
|
|
||||||
#### 9c — Sanitizer
|
|
||||||
- [x] `app/api/middleware/sanitizer.py`:
|
|
||||||
- Response middleware that scans response bodies
|
|
||||||
- Strips: system prompt fragments, agent internal reasoning, tool schemas, routing metadata
|
|
||||||
- Pattern-based detection + exact match against known prompt fingerprints
|
|
||||||
- Logs sanitization events for monitoring
|
|
||||||
|
|
||||||
- **Outcome:** Secure, rate-limited API with prompt IP protection.
|
|
||||||
|
|
||||||
### Step 10 — Plugin Marketplace ✅
|
|
||||||
- [x] `app/marketplace/plugin_registry.py`:
|
|
||||||
- `PluginRegistry`:
|
|
||||||
- `async list_plugins(category, query, page, sort) -> PluginListResponse`
|
|
||||||
- `async get_plugin(plugin_id) -> PluginManifest | None`
|
|
||||||
- `async submit_plugin(manifest: PluginManifest, package_s3_key: str) -> str` — returns plugin_id, sets status = 'pending_review'
|
|
||||||
- `async approve_plugin(plugin_id) -> None` — admin only, sets status = 'approved'
|
|
||||||
- `async reject_plugin(plugin_id, reason: str) -> None`
|
|
||||||
- [x] `app/marketplace/plugin_review.py`:
|
|
||||||
- `ReviewQueue`:
|
|
||||||
- `async get_pending() -> list[dict]`
|
|
||||||
- `async submit_review(plugin_id, reviewer_id, decision, notes) -> None`
|
|
||||||
- Security checklist enforced before approval: manifest schema valid, permissions are from allowed set, no binary blobs in manifest
|
|
||||||
- [x] `app/marketplace/revenue_share.py`:
|
|
||||||
- `RevenueShare`:
|
|
||||||
- `async record_install(plugin_id, user_id, amount_cents) -> None`
|
|
||||||
- `async payout_developer(plugin_id, period) -> None` — Stripe Connect transfer: 70% to developer
|
|
||||||
- `async get_earnings(developer_id, period) -> dict`
|
|
||||||
- **Outcome:** Plugin marketplace with catalog, review workflow, and revenue split.
|
|
||||||
|
|
||||||
### Step 11 — Billing & Tier management ✅
|
|
||||||
- [x] `app/billing/stripe_service.py`:
|
|
||||||
- `create_checkout_session(user_id, tier) -> str`
|
|
||||||
- `handle_webhook(payload, sig_header) -> None`: processes `checkout.session.completed`, `customer.subscription.updated`, `customer.subscription.deleted`, `invoice.payment_failed`
|
|
||||||
- `get_subscription(user_id) -> dict | None`
|
|
||||||
- `cancel_subscription(user_id) -> None`
|
|
||||||
- [x] `app/billing/tier_manager.py`:
|
|
||||||
- `TierManager`:
|
|
||||||
- Feature matrix:
|
|
||||||
```python
|
|
||||||
FEATURES = {
|
|
||||||
'free': {
|
|
||||||
'agents': 3,
|
|
||||||
'batch_active': 2,
|
|
||||||
'cloud_storage_gb': 0,
|
|
||||||
'backup_gb': 0,
|
|
||||||
'providers': 1,
|
|
||||||
'batch_builder': False,
|
|
||||||
'plugin_marketplace': False,
|
|
||||||
'sso': False,
|
|
||||||
},
|
|
||||||
'pro': {
|
|
||||||
'agents': -1, # unlimited
|
|
||||||
'batch_active': 10,
|
|
||||||
'cloud_storage_gb': 5,
|
|
||||||
'backup_gb': 5,
|
|
||||||
'providers': -1,
|
|
||||||
'batch_builder': False,
|
|
||||||
'plugin_marketplace': False,
|
|
||||||
'sso': False,
|
|
||||||
},
|
|
||||||
'power': {
|
|
||||||
'agents': -1,
|
|
||||||
'batch_active': -1, # unlimited
|
|
||||||
'cloud_storage_gb': 25,
|
|
||||||
'backup_gb': 25,
|
|
||||||
'providers': -1,
|
|
||||||
'batch_builder': True,
|
|
||||||
'plugin_marketplace': True,
|
|
||||||
'sso': False,
|
|
||||||
},
|
|
||||||
'team': {
|
|
||||||
'agents': -1,
|
|
||||||
'batch_active': -1,
|
|
||||||
'cloud_storage_gb': -1,
|
|
||||||
'backup_gb': -1,
|
|
||||||
'providers': -1,
|
|
||||||
'batch_builder': True,
|
|
||||||
'plugin_marketplace': True,
|
|
||||||
'sso': True,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
```
|
|
||||||
- `get_tier(user_id) -> BillingTier`
|
|
||||||
- `check_feature(user_id, feature) -> bool`
|
|
||||||
- `get_rate_limit(tier) -> int`
|
|
||||||
- `check_quota(user_id) -> bool` — checks cloud_storage_gb current usage vs limit
|
|
||||||
- [x] `app/billing/__init__.py`: exports `stripe_service` and `tier_manager` singletons
|
|
||||||
- [x] `app/api/routes/billing.py`: refactored to delegate to `StripeService`
|
|
||||||
- [x] `app/api/routes/storage.py` and `backup.py`: `_check_quota` now delegates to `tier_manager.enforce_quota` / `enforce_backup_quota`
|
|
||||||
- **Outcome:** Stripe integration with tier-based feature gating matching Free/Pro(15€)/Power(29€)/Team(49€/seat).
|
|
||||||
|
|
||||||
### Step 12 — Database (auth/billing/marketplace only)
|
|
||||||
- [x] PostgreSQL schema via Alembic:
|
|
||||||
- `users`: `id UUID PK`, `email UNIQUE`, `password_hash`, `tier` (default 'free'), `stripe_customer_id`, `created_at`, `updated_at`
|
|
||||||
- `refresh_tokens`: `id UUID PK`, `user_id FK`, `token_hash`, `expires_at`, `created_at`
|
|
||||||
- `subscriptions`: `id UUID PK`, `user_id FK`, `stripe_subscription_id`, `tier`, `status`, `current_period_end`, `created_at`
|
|
||||||
- `backup_metadata`: `id UUID PK`, `user_id FK`, `s3_key`, `version`, `timestamp`, `checksum`, `size_bytes`, `created_at`
|
|
||||||
- `storage_records`: `id UUID PK`, `user_id FK`, `table_name VARCHAR`, `s3_key`, `checksum`, `size_bytes`, `created_at`, `updated_at` — metadata only, no plaintext
|
|
||||||
- `plugins`: `id UUID PK`, `name`, `description`, `version`, `author_id FK`, `category`, `status` (pending_review/approved/rejected), `price_cents`, `s3_package_key`, `install_count`, `avg_rating`, `created_at`
|
|
||||||
- `plugin_installations`: `id UUID PK`, `plugin_id FK`, `user_id FK`, `installed_at`
|
|
||||||
- `plugin_reviews`: `id UUID PK`, `plugin_id FK`, `reviewer_id FK`, `decision`, `notes`, `reviewed_at`
|
|
||||||
- `revenue_events`: `id UUID PK`, `plugin_id FK`, `user_id FK`, `amount_cents`, `developer_share_cents`, `stripe_transfer_id`, `created_at`
|
|
||||||
- [x] Initial Alembic migration
|
|
||||||
- [x] SQLAlchemy models in `app/models.py`
|
|
||||||
- **Outcome:** Auth, billing, storage metadata, and marketplace persistence. Zero user data in plaintext.
|
|
||||||
|
|
||||||
### Step 13 — Testing & deployment ✅
|
|
||||||
- [x] `tests/conftest.py`: TestClient fixture, mock LLM fixture (`AsyncMock` returning canned responses), mock agent fixture, test DB (SQLite in-memory for speed), mock S3 (moto), mock Pinecone
|
|
||||||
- [x] `tests/test_orchestrator.py`: classify_intent routing, single agent, pipeline, plan mode
|
|
||||||
- [x] `tests/test_agents.py`: each agent with mocked tools
|
|
||||||
- [x] `tests/test_auth.py`: register → login → access protected → refresh → expired token
|
|
||||||
- [x] `tests/test_backup.py`: upload → download → history → delete, tier limit enforcement
|
|
||||||
- [x] `tests/test_storage.py`: create record → list → download → update → delete, checksum rejection, quota enforcement
|
|
||||||
- [x] `tests/test_plugins.py`: list plugins, install, uninstall, revenue event creation, tier gate (free user blocked)
|
|
||||||
- [x] `Dockerfile` optimized for production (gunicorn + uvicorn workers)
|
|
||||||
- [x] GitHub Actions CI: lint (ruff), test (pytest), build Docker image
|
|
||||||
- **Outcome:** Fully tested, deployable backend.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## API Contract Summary
|
|
||||||
|
|
||||||
| Method | Endpoint | Auth | Request | Response |
|
|
||||||
|--------|----------|------|---------|----------|
|
|
||||||
| POST | `/api/v1/auth/register` | No | `{email, password}` | `AuthTokens` |
|
|
||||||
| POST | `/api/v1/auth/login` | No | `{email, password}` | `AuthTokens` |
|
|
||||||
| POST | `/api/v1/auth/refresh` | No | `{refresh_token}` | `AuthTokens` |
|
|
||||||
| GET | `/api/v1/auth/me` | JWT | — | `UserProfile` |
|
|
||||||
| POST | `/api/v1/chat` | JWT | `ChatRequest` | `ChatResponse \| ExecutionPlan` |
|
|
||||||
| WS | `/api/v1/chat/stream` | JWT | `ChatRequest` (first frame) | Token stream + final JSON |
|
|
||||||
| GET | `/api/v1/plans/playbook` | JWT | — | `ExecutionPlan[]` |
|
|
||||||
| GET | `/api/v1/plans/playbook/:id` | JWT | — | `ExecutionPlan` |
|
|
||||||
| POST | `/api/v1/storage/records` | JWT | `StorageRecordCreate` | `{id, created_at}` |
|
|
||||||
| GET | `/api/v1/storage/records` | JWT | `?table&page&limit` | `RecordMeta[]` |
|
|
||||||
| GET | `/api/v1/storage/records/:id` | JWT | — | Binary blob |
|
|
||||||
| PUT | `/api/v1/storage/records/:id` | JWT | `StorageRecordUpdate` | `{ok: true}` |
|
|
||||||
| DELETE | `/api/v1/storage/records/:id` | JWT | — | `{ok: true}` |
|
|
||||||
| POST | `/api/v1/storage/vectors/upsert` | JWT | `VectorUpsertRequest` | `{upserted: int}` |
|
|
||||||
| POST | `/api/v1/storage/vectors/search` | JWT | `VectorSearchRequest` | `VectorSearchResponse` |
|
|
||||||
| DELETE | `/api/v1/storage/vectors` | JWT | `{ids: list[str]}` | `{ok: true}` |
|
|
||||||
| PUT | `/api/v1/backup` | JWT | Binary blob + headers | `{ok: true}` |
|
|
||||||
| GET | `/api/v1/backup` | JWT | — | Binary blob |
|
|
||||||
| GET | `/api/v1/backup/history` | JWT | — | `BackupMetadata[]` |
|
|
||||||
| DELETE | `/api/v1/backup/:id` | JWT | — | `{ok: true}` |
|
|
||||||
| GET | `/api/v1/plugins` | JWT | `?category&q&page&sort` | `PluginListResponse` |
|
|
||||||
| GET | `/api/v1/plugins/:id` | JWT | — | `PluginManifest` + stats |
|
|
||||||
| POST | `/api/v1/plugins/:id/install` | JWT | `PluginInstallRequest` | `{ok, download_url}` |
|
|
||||||
| DELETE | `/api/v1/plugins/:id/install` | JWT | — | `{ok: true}` |
|
|
||||||
| POST | `/api/v1/billing/checkout` | JWT | `{tier}` | `{checkout_url}` |
|
|
||||||
| POST | `/api/v1/billing/webhook` | Stripe sig | Stripe event | `{ok: true}` |
|
|
||||||
| GET | `/api/v1/billing/subscription` | JWT | — | Subscription info |
|
|
||||||
| DELETE | `/api/v1/billing/subscription` | JWT | — | `{ok: true}` |
|
|
||||||
| GET | `/api/v1/health` | No | — | `{status, version}` |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Stack
|
|
||||||
|
|
||||||
| Layer | Technology |
|
|
||||||
|-------|-----------|
|
|
||||||
| Framework | FastAPI + Uvicorn |
|
|
||||||
| LLM | LangChain + langchain-openai |
|
|
||||||
| Auth | PyJWT + bcrypt + OAuth2 |
|
|
||||||
| Billing | stripe-python + Stripe Connect |
|
|
||||||
| Blob storage | boto3 (S3) |
|
|
||||||
| Vector store | Pinecone or Qdrant (configurable) |
|
|
||||||
| Database | PostgreSQL + SQLAlchemy + Alembic |
|
|
||||||
| Rate limiting | slowapi |
|
|
||||||
| Testing | pytest + pytest-asyncio + httpx + moto (S3 mock) |
|
|
||||||
| Deployment | Docker → fly.io / Railway / AWS ECS |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Development Rules
|
|
||||||
|
|
||||||
1. **NEVER persist user data in plaintext.** The DB stores only auth, billing, storage metadata, and marketplace data. User context arrives in requests and is discarded. Cloud blobs are E2E encrypted client-side — backend only stores opaque bytes.
|
|
||||||
2. **NEVER expose prompts.** System prompts are composed server-side from fragments. Responses are sanitized before sending. In plan mode, `prompt_template` fields are reference IDs only.
|
|
||||||
3. **NEVER decrypt user blobs.** `app/storage/encryption.py` only verifies checksums. No decryption key ever reaches the backend.
|
|
||||||
4. **Stateless request handling.** No server-side session state. All context comes from the client + JWT.
|
|
||||||
5. **Type hints everywhere.** All functions have full type annotations.
|
|
||||||
6. **Test every agent.** Each chat agent has unit tests with mocked LLM responses.
|
|
||||||
7. **Structured logging.** JSON logs with request ID correlation.
|
|
||||||
8. **Tier gates are enforced server-side.** Never trust client-reported tier. Always fetch from DB via `TierManager.get_tier(user_id)`.
|
|
||||||
9. **One step at a time.** Implement one numbered step per session. When the step is fully done, mark all its checkboxes as `[x]` in this file and commit with message `step N complete: <outcome line>`.
|
|
||||||
794
README.md
794
README.md
@@ -1,793 +1,5 @@
|
|||||||
# Adiuva Cloud API
|
## DEV
|
||||||
|
Run in DEV with command:
|
||||||
**AI-powered project management backend with E2E encrypted cloud storage, LLM orchestration, and a plugin marketplace.**
|
|
||||||
|
|
||||||
Built with FastAPI · Python 3.12 · PostgreSQL · LangChain · Stripe · AWS S3
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Table of Contents
|
|
||||||
|
|
||||||
- [Overview](#overview)
|
|
||||||
- [Architecture](#architecture)
|
|
||||||
- [Key Features](#key-features)
|
|
||||||
- [Tech Stack](#tech-stack)
|
|
||||||
- [Getting Started](#getting-started)
|
|
||||||
- [Docker Deployment](#docker-deployment)
|
|
||||||
- [Environment Variables](#environment-variables)
|
|
||||||
- [API Reference](#api-reference)
|
|
||||||
- [Data Model](#data-model)
|
|
||||||
- [AI Agent System](#ai-agent-system)
|
|
||||||
- [Orchestration & Execution Plans](#orchestration--execution-plans)
|
|
||||||
- [Middleware](#middleware)
|
|
||||||
- [Storage Layer](#storage-layer)
|
|
||||||
- [Billing & Tiers](#billing--tiers)
|
|
||||||
- [Plugin Marketplace](#plugin-marketplace)
|
|
||||||
- [Testing](#testing)
|
|
||||||
- [Project Structure](#project-structure)
|
|
||||||
- [License](#license)
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Overview
|
|
||||||
|
|
||||||
Adiuva Cloud API is the FastAPI backend that powers the **Adiuva Electron desktop app**. It provides LLM-powered chat orchestration, end-to-end encrypted cloud storage, a vector search engine, an encrypted backup system, a plugin marketplace with revenue sharing, and Stripe-based subscription billing across four tiers.
|
|
||||||
|
|
||||||
### Design Principles
|
|
||||||
|
|
||||||
1. **Never persist user data in plaintext** — the database stores only auth, billing, storage metadata, and marketplace data. All user content is E2E encrypted by the client before reaching the server.
|
|
||||||
2. **Never expose prompts** — system prompts stay server-side; responses are sanitized to strip any leaked prompt fragments.
|
|
||||||
3. **Never decrypt user blobs** — the backend performs only checksum verification; no decryption keys ever reach the server.
|
|
||||||
4. **Stateless request handling** — all context comes from the client and JWT; no server-side session state.
|
|
||||||
5. **Tier gates enforced server-side** — the server always reads the current tier from the database, never trusting client-reported values.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Architecture
|
|
||||||
|
|
||||||
```
|
```
|
||||||
┌──────────────┐ ┌────────────────────────────────────────────────────────┐
|
uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload --log-config logging.conf
|
||||||
│ Electron │ │ FastAPI (Uvicorn / Gunicorn) │
|
|
||||||
│ Desktop App │────▶│ │
|
|
||||||
│ (Client) │◀────│ Middleware: RateLimit → Sanitizer → CORS → Router │
|
|
||||||
└──────────────┘ │ │
|
|
||||||
│ ┌──────────────────┐ ┌────────────────────────────┐ │
|
|
||||||
│ │ Auth Routes │ │ Chat Routes │ │
|
|
||||||
│ │ Billing Routes │ │ ↓ │ │
|
|
||||||
│ │ Storage Routes │ │ Orchestrator (GPT-4o-mini)│ │
|
|
||||||
│ │ Backup Routes │ │ ↓ classify intent │ │
|
|
||||||
│ │ Plugin Routes │ │ Agent Registry │ │
|
|
||||||
│ │ Vector Routes │ │ ↓ │ │
|
|
||||||
│ │ Plans Routes │ │ TaskAgent | ProjectAgent │ │
|
|
||||||
│ └──────────────────┘ │ NoteAgent | CheckptAgent │ │
|
|
||||||
│ │ (GPT-4o + LangChain) │ │
|
|
||||||
│ └────────────────────────────┘ │
|
|
||||||
└────────────────────────────────────────────────────────┘
|
|
||||||
│ │ │
|
|
||||||
┌────────▼───┐ ┌───────▼───────┐ ┌──▼─────────────┐
|
|
||||||
│ PostgreSQL │ │ AWS S3 │ │ Pinecone / │
|
|
||||||
│ (Auth, │ │ (E2E blobs, │ │ Qdrant │
|
|
||||||
│ Billing, │ │ backups) │ │ (Vectors) │
|
|
||||||
│ Metadata) │ └───────────────┘ └────────────────┘
|
|
||||||
└────────────┘
|
|
||||||
│
|
|
||||||
┌────────▼───┐
|
|
||||||
│ Stripe │
|
|
||||||
│ (Billing, │
|
|
||||||
│ Connect) │
|
|
||||||
└────────────┘
|
|
||||||
```
|
```
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Key Features
|
|
||||||
|
|
||||||
1. **LLM-powered orchestration** — GPT-4o-mini classifies user intent and routes to the appropriate domain agent.
|
|
||||||
2. **4 specialized AI agents** — Tasks (8 tools), Projects (6 tools), Checkpoints (4 tools), Notes (5 tools), all powered by GPT-4o via LangChain.
|
|
||||||
3. **Execution plans & playbooks** — Server-side prompt template registry; clients receive only opaque template IDs, never raw prompts.
|
|
||||||
4. **E2E encrypted cloud storage** — The backend never decrypts user data; SHA-256 checksum verification uses constant-time comparison to prevent timing attacks.
|
|
||||||
5. **Cloud vector store** — Pinecone or Qdrant with user-isolated namespaces and encrypted blob payloads.
|
|
||||||
6. **Encrypted backup system** — Tiered storage limits with `If-Modified-Since` support for efficient syncing.
|
|
||||||
7. **Plugin marketplace** — Catalog, admin review/approval workflow, security checklist, and 70/30 revenue sharing via Stripe Connect.
|
|
||||||
8. **Stripe billing** — Four-tier subscription model (Free / Pro / Power / Team) with checkout sessions and full webhook lifecycle handling.
|
|
||||||
9. **JWT authentication** — Access + refresh tokens with bcrypt password hashing, SHA-256 token hashing, and automatic rotation.
|
|
||||||
10. **Prompt IP protection** — Sanitizer middleware strips system prompts, reasoning markers, tool schemas, and agent routing metadata from all chat responses.
|
|
||||||
11. **Tier-based rate limiting** — Sliding-window per-user limiter scaling from 20 to 200 requests/min by subscription tier.
|
|
||||||
12. **Zero-trust data model** — User content is never stored in plaintext; the database holds only authentication, billing, and metadata records.
|
|
||||||
13. **WebSocket streaming** — Real-time chat with 30-second heartbeat keep-alive and chunked text delivery.
|
|
||||||
14. **Alembic migrations** — Versioned schema management with seed data for the plugin marketplace.
|
|
||||||
15. **Comprehensive test suite** — In-memory SQLite + moto S3 mocks, per-tier test fixtures, and full API coverage without external dependencies.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Tech Stack
|
|
||||||
|
|
||||||
| Package | Version | Purpose |
|
|
||||||
|---|---|---|
|
|
||||||
| `fastapi` | ≥ 0.115.0 | Web framework |
|
|
||||||
| `uvicorn[standard]` | ≥ 0.34.0 | ASGI development server |
|
|
||||||
| `gunicorn` | ≥ 22.0.0 | Production process manager |
|
|
||||||
| `langchain` | ≥ 0.3.0 | LLM orchestration framework |
|
|
||||||
| `langchain-openai` | ≥ 0.3.0 | OpenAI LLM provider integration |
|
|
||||||
| `litellm` | ≥ 1.50.0 | Universal LLM gateway (100+ providers) |
|
|
||||||
| `pydantic` | ≥ 2.10.0 | Data validation and serialization |
|
|
||||||
| `pydantic-settings` | ≥ 2.7.0 | Environment-based configuration |
|
|
||||||
| `python-jose[cryptography]` | ≥ 3.3.0 | JWT encoding and decoding |
|
|
||||||
| `stripe` | ≥ 11.0.0 | Billing and payment integration |
|
|
||||||
| `boto3` | ≥ 1.35.0 | AWS S3 client |
|
|
||||||
| `slowapi` | ≥ 0.1.9 | Rate limiting utilities |
|
|
||||||
| `sqlalchemy` | ≥ 2.0.0 | Async ORM and query builder |
|
|
||||||
| `asyncpg` | ≥ 0.30.0 | PostgreSQL async driver |
|
|
||||||
| `alembic` | ≥ 1.14.0 | Database migration management |
|
|
||||||
| `bcrypt` | ≥ 4.2.0 | Password hashing |
|
|
||||||
| `python-dotenv` | ≥ 1.0.0 | `.env` file loading |
|
|
||||||
| `httpx` | ≥ 0.28.0 | Async HTTP client (used in tests) |
|
|
||||||
| `websockets` | ≥ 14.0 | WebSocket protocol support |
|
|
||||||
| `psycopg2-binary` | ≥ 2.9.0 | Synchronous PostgreSQL driver (Alembic) |
|
|
||||||
| `pinecone` | ≥ 5.0.0 | Pinecone vector store client |
|
|
||||||
| `qdrant-client` | ≥ 1.7.0 | Qdrant vector store client |
|
|
||||||
| `pytest` | ≥ 8.0.0 | Test framework |
|
|
||||||
| `pytest-asyncio` | ≥ 0.24.0 | Async test support |
|
|
||||||
| `aiosqlite` | ≥ 0.20.0 | In-memory SQLite for tests |
|
|
||||||
| `moto[s3]` | ≥ 5.0.0 | AWS S3 mock for tests |
|
|
||||||
| `ruff` | ≥ 0.8.0 | Linter and formatter |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Getting Started
|
|
||||||
|
|
||||||
### Prerequisites
|
|
||||||
|
|
||||||
- Python 3.12+
|
|
||||||
- PostgreSQL 16+
|
|
||||||
- An OpenAI API key (for LLM features)
|
|
||||||
- Stripe API keys (optional — billing stubs gracefully when unconfigured)
|
|
||||||
- AWS credentials (optional — needed for S3 storage in production)
|
|
||||||
|
|
||||||
### Installation
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Clone the repository
|
|
||||||
git clone <repo-url> && cd adiuva-api
|
|
||||||
|
|
||||||
# Create a virtual environment
|
|
||||||
python -m venv .venv && source .venv/bin/activate
|
|
||||||
|
|
||||||
# Install dependencies
|
|
||||||
pip install -r requirements.txt
|
|
||||||
|
|
||||||
# Configure environment
|
|
||||||
cp .env.example .env
|
|
||||||
# Edit .env with your DATABASE_URL, OPENAI_API_KEY, etc.
|
|
||||||
```
|
|
||||||
|
|
||||||
### Database Setup
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Start PostgreSQL (or use the Docker Compose database)
|
|
||||||
docker compose up db -d
|
|
||||||
|
|
||||||
# Run migrations
|
|
||||||
alembic upgrade head
|
|
||||||
```
|
|
||||||
|
|
||||||
### Run the Development Server
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uvicorn app.main:app --reload --host 0.0.0.0 --port 8000
|
|
||||||
```
|
|
||||||
|
|
||||||
Interactive API docs are available at [http://localhost:8000/docs](http://localhost:8000/docs) in development mode (`ENV=dev`). The `/docs` endpoint is disabled in production.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Docker Deployment
|
|
||||||
|
|
||||||
### Quick Start
|
|
||||||
|
|
||||||
```bash
|
|
||||||
docker compose up --build
|
|
||||||
```
|
|
||||||
|
|
||||||
This starts two services:
|
|
||||||
|
|
||||||
- **app** — FastAPI server on port `8000`
|
|
||||||
- **db** — PostgreSQL 16 (Alpine) on port `5432` with a persistent volume and health checks
|
|
||||||
|
|
||||||
The compose file also includes optional services for fully local deployments:
|
|
||||||
|
|
||||||
- **minio** — S3-compatible object storage on ports `9000` (API) and `9001` (console)
|
|
||||||
- **qdrant** — Vector search engine on ports `6333` (HTTP) and `6334` (gRPC)
|
|
||||||
|
|
||||||
### Dockerfile Details
|
|
||||||
|
|
||||||
The Dockerfile uses a multi-stage build:
|
|
||||||
|
|
||||||
1. **Builder stage** — Installs Python dependencies into a virtual environment.
|
|
||||||
2. **Runtime stage** — Copies only the venv, app source, and Alembic migrations. Runs as a non-root user (`appuser`).
|
|
||||||
3. **Production server** — Gunicorn with 4 Uvicorn workers, 120-second timeout, listening on port 8000.
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Production command (run by the container)
|
|
||||||
gunicorn app.main:app -k uvicorn.workers.UvicornWorker -w 4 --timeout 120 -b 0.0.0.0:8000
|
|
||||||
```
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Homelab / Self-Hosted Deployment
|
|
||||||
|
|
||||||
You can run the entire stack locally on a homelab with **no cloud dependencies except the LLM provider**. The compose file includes MinIO (S3 replacement) and Qdrant (vector store) out of the box.
|
|
||||||
|
|
||||||
### 1. Start all services
|
|
||||||
|
|
||||||
```bash
|
|
||||||
docker compose up -d
|
|
||||||
```
|
|
||||||
|
|
||||||
This starts PostgreSQL, MinIO, and Qdrant alongside the app.
|
|
||||||
|
|
||||||
### 2. Create the MinIO bucket
|
|
||||||
|
|
||||||
Open the MinIO console at [http://localhost:9001](http://localhost:9001) (login: `minioadmin` / `minioadmin`) and create a bucket named `adiuva`, or use the CLI:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
docker compose exec minio mc alias set local http://localhost:9000 minioadmin minioadmin
|
|
||||||
docker compose exec minio mc mb local/adiuva
|
|
||||||
```
|
|
||||||
|
|
||||||
### 3. Configure your `.env`
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Database (uses the compose PostgreSQL)
|
|
||||||
DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/adiuva
|
|
||||||
|
|
||||||
# S3 → MinIO
|
|
||||||
S3_BUCKET=adiuva
|
|
||||||
S3_REGION=us-east-1
|
|
||||||
S3_ENDPOINT_URL=http://minio:9000
|
|
||||||
AWS_ACCESS_KEY_ID=minioadmin
|
|
||||||
AWS_SECRET_ACCESS_KEY=minioadmin
|
|
||||||
|
|
||||||
# Vector store → local Qdrant (leave PINECONE_API_KEY empty)
|
|
||||||
QDRANT_URL=http://qdrant:6333
|
|
||||||
QDRANT_API_KEY=
|
|
||||||
PINECONE_API_KEY=
|
|
||||||
|
|
||||||
# Billing — leave empty to stub (no Stripe needed)
|
|
||||||
STRIPE_SECRET_KEY=
|
|
||||||
STRIPE_WEBHOOK_SECRET=
|
|
||||||
|
|
||||||
# LLM — the only external service
|
|
||||||
OPENAI_API_KEY=sk-...
|
|
||||||
LLM_MODEL=gpt-4o
|
|
||||||
LLM_ROUTER_MODEL=gpt-4o-mini
|
|
||||||
|
|
||||||
# Auth
|
|
||||||
JWT_SECRET=your-secret-here
|
|
||||||
ENV=dev
|
|
||||||
```
|
|
||||||
|
|
||||||
### 4. Run migrations
|
|
||||||
|
|
||||||
```bash
|
|
||||||
docker compose exec app alembic upgrade head
|
|
||||||
```
|
|
||||||
|
|
||||||
### What runs where
|
|
||||||
|
|
||||||
| Service | Runs on | Port | Notes |
|
|
||||||
|---|---|---|---|
|
|
||||||
| FastAPI app | Docker | 8000 | API server |
|
|
||||||
| PostgreSQL | Docker | 5432 | Auth, billing, metadata |
|
|
||||||
| MinIO | Docker | 9000 / 9001 | S3-compatible blob & backup storage |
|
|
||||||
| Qdrant | Docker | 6333 / 6334 | Vector search (replaces Pinecone) |
|
|
||||||
| Stripe | — | — | Stubbed when keys are empty |
|
|
||||||
| OpenAI / LLM | Cloud | — | Only external dependency |
|
|
||||||
|
|
||||||
> **Want fully offline AI too?** Set `LLM_MODEL=ollama/llama3` and `LLM_ROUTER_MODEL=ollama/llama3`, then add an Ollama container or point at a local Ollama instance. See the [LLM provider switching](#switching-llm-providers) section.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Environment Variables
|
|
||||||
|
|
||||||
All variables are loaded from a `.env` file via Pydantic Settings. Source: `app/config/settings.py`
|
|
||||||
|
|
||||||
| Variable | Type | Default | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `DATABASE_URL` | `str` | `postgresql+asyncpg://postgres:postgres@localhost:5432/adiuva` | Async SQLAlchemy connection string |
|
|
||||||
| `JWT_SECRET` | `str` | `change-me-in-production` | HMAC secret for JWT signing |
|
|
||||||
| `JWT_ALGORITHM` | `str` | `HS256` | JWT signing algorithm |
|
|
||||||
| `JWT_ACCESS_TOKEN_EXPIRE_MINUTES` | `int` | `30` | Access token time-to-live |
|
|
||||||
| `JWT_REFRESH_TOKEN_EXPIRE_DAYS` | `int` | `30` | Refresh token time-to-live |
|
|
||||||
| `STRIPE_SECRET_KEY` | `str` | `""` | Stripe API key (empty = stub mode) |
|
|
||||||
| `STRIPE_WEBHOOK_SECRET` | `str` | `""` | Stripe webhook signature secret |
|
|
||||||
| `S3_BUCKET` | `str` | `""` | S3 bucket for encrypted blobs and backups |
|
|
||||||
| `S3_REGION` | `str` | `us-east-1` | AWS region |
|
|
||||||
| `S3_ENDPOINT_URL` | `str` | `""` | Custom S3 endpoint (e.g. `http://minio:9000` for MinIO). Leave empty for AWS. |
|
|
||||||
| `AWS_ACCESS_KEY_ID` | `str` | `""` | AWS credentials |
|
|
||||||
| `AWS_SECRET_ACCESS_KEY` | `str` | `""` | AWS credentials |
|
|
||||||
| `PINECONE_API_KEY` | `str` | `""` | Pinecone API key (if set, Pinecone is used for vectors) |
|
|
||||||
| `PINECONE_INDEX` | `str` | `adiuva` | Pinecone index name |
|
|
||||||
| `QDRANT_URL` | `str` | `""` | Qdrant URL (used when Pinecone is not configured) |
|
|
||||||
| `QDRANT_API_KEY` | `str` | `""` | Qdrant API key |
|
|
||||||
| `OPENAI_API_KEY` | `str` | `""` | OpenAI key for LLM agent calls |
|
|
||||||
| `LLM_MODEL` | `str` | `gpt-4o` | LiteLLM model identifier for agents (e.g. `anthropic/claude-3.5-sonnet`, `gemini/gemini-pro`, `ollama/llama3`) |
|
|
||||||
| `LLM_ROUTER_MODEL` | `str` | `gpt-4o-mini` | Lighter model used for intent classification / routing |
|
|
||||||
| `CORS_ORIGINS` | `list[str]` | `["app://.", "http://localhost:3000", "http://localhost:5173"]` | Allowed CORS origins |
|
|
||||||
| `ENV` | `Literal` | `dev` | `dev` or `prod` — controls `/docs` visibility and SQL echo |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## API Reference
|
|
||||||
|
|
||||||
All routes are prefixed with `/api/v1`. **27 endpoints** total (25 REST + 1 WebSocket + 1 health check).
|
|
||||||
|
|
||||||
### Health
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `GET` | `/api/v1/health` | No | Returns `{"status": "ok", "version": "0.1.0"}` |
|
|
||||||
|
|
||||||
### Auth
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `POST` | `/api/v1/auth/register` | No | Create account with bcrypt-hashed password, returns `AuthTokens` |
|
|
||||||
| `POST` | `/api/v1/auth/login` | No | Validate credentials, returns `AuthTokens` |
|
|
||||||
| `POST` | `/api/v1/auth/refresh` | No | Rotate refresh token, returns new `AuthTokens` |
|
|
||||||
| `GET` | `/api/v1/auth/me` | JWT | Returns `UserProfile` for the authenticated user |
|
|
||||||
|
|
||||||
### Chat
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `POST` | `/api/v1/chat` | JWT | Route message through the orchestrator; returns `ChatResponse` or `ExecutionPlan` depending on execution mode |
|
|
||||||
| `WS` | `/api/v1/chat/stream` | JWT (query param `?token=`) | Streaming chat — first frame is a `ChatRequest`, server yields text chunks, final frame is `{"done": true, "response": "...", "actions": [...]}`. 30-second heartbeat ping. |
|
|
||||||
|
|
||||||
### Plans
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `GET` | `/api/v1/plans/playbook` | JWT | List all cached execution plan playbooks |
|
|
||||||
| `GET` | `/api/v1/plans/playbook/{plan_id}` | JWT | Retrieve a specific playbook by ID |
|
|
||||||
|
|
||||||
### Storage (Cloud Records)
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `POST` | `/api/v1/storage/records` | JWT | Upload an E2E encrypted record (verifies checksum, enforces storage quota) |
|
|
||||||
| `GET` | `/api/v1/storage/records` | JWT | List record metadata with pagination (`?table`, `?page`, `?limit`); no blob bytes returned |
|
|
||||||
| `GET` | `/api/v1/storage/records/{id}` | JWT | Download encrypted blob with `X-Checksum` response header |
|
|
||||||
| `PUT` | `/api/v1/storage/records/{id}` | JWT | Replace an existing blob (verifies checksum, enforces quota) |
|
|
||||||
| `DELETE` | `/api/v1/storage/records/{id}` | JWT | Delete a record and its S3 blob |
|
|
||||||
|
|
||||||
### Vectors (Cloud Vector Store)
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `POST` | `/api/v1/storage/vectors/upsert` | JWT | Verify checksums and upsert encrypted vectors |
|
|
||||||
| `POST` | `/api/v1/storage/vectors/search` | JWT | Search user-scoped vector namespace |
|
|
||||||
| `DELETE` | `/api/v1/storage/vectors` | JWT | Delete vectors by ID list |
|
|
||||||
|
|
||||||
### Backup
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `PUT` | `/api/v1/backup` | JWT | Upload encrypted backup blob with custom headers (`X-Backup-Version`, `X-Backup-Timestamp`, `X-Backup-Checksum`). Tier quota enforced. |
|
|
||||||
| `GET` | `/api/v1/backup` | JWT | Download latest backup blob. Supports `If-Modified-Since`. |
|
|
||||||
| `GET` | `/api/v1/backup/history` | JWT | List backup metadata (no blob content) |
|
|
||||||
| `DELETE` | `/api/v1/backup/{backup_id}` | JWT | Delete a specific backup |
|
|
||||||
|
|
||||||
### Plugins (Marketplace)
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `GET` | `/api/v1/plugins` | JWT (Power+) | Browse the marketplace (`?category`, `?q`, `?page`, `?sort=rating\|installs\|newest`) |
|
|
||||||
| `GET` | `/api/v1/plugins/{id}` | JWT (Power+) | Plugin detail with install count and ratings |
|
|
||||||
| `POST` | `/api/v1/plugins/{id}/install` | JWT (Power+) | Install plugin; triggers Stripe Connect revenue split for paid plugins |
|
|
||||||
| `DELETE` | `/api/v1/plugins/{id}/install` | JWT | Uninstall plugin |
|
|
||||||
|
|
||||||
### Billing
|
|
||||||
|
|
||||||
| Method | Path | Auth | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `POST` | `/api/v1/billing/checkout` | JWT | Create a Stripe checkout session, returns `{"checkout_url": "..."}` |
|
|
||||||
| `POST` | `/api/v1/billing/webhook` | Stripe signature | Handle Stripe events: `checkout.session.completed`, `customer.subscription.updated`, `customer.subscription.deleted`, `invoice.payment_failed` |
|
|
||||||
| `GET` | `/api/v1/billing/subscription` | JWT | Get current subscription information |
|
|
||||||
| `DELETE` | `/api/v1/billing/subscription` | JWT | Cancel subscription and revert to free tier |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Data Model
|
|
||||||
|
|
||||||
9 tables managed by Alembic migrations. Source: `app/models.py`
|
|
||||||
|
|
||||||
### Tables
|
|
||||||
|
|
||||||
| Table | Primary Key | Key Columns | Purpose |
|
|
||||||
|---|---|---|---|
|
|
||||||
| `users` | `id` (UUID) | `email` (unique), `password_hash`, `tier`, `stripe_customer_id`, timestamps | User accounts |
|
|
||||||
| `refresh_tokens` | `id` (UUID) | `user_id` (FK), `token_hash` (SHA-256, unique), `expires_at` | Hashed refresh tokens for rotation |
|
|
||||||
| `subscriptions` | `id` (UUID) | `user_id` (FK, unique), `stripe_subscription_id`, `tier`, `status`, `current_period_end` | Stripe subscription records |
|
|
||||||
| `storage_records` | `id` (UUID) | `user_id` (FK), `table_name`, `s3_key`, `checksum`, `size_bytes`, timestamps | S3 blob metadata (no plaintext content) |
|
|
||||||
| `backup_metadata` | `id` (UUID) | `user_id` (FK), `s3_key`, `version`, `timestamp`, `checksum`, `size_bytes` | Backup manifests |
|
|
||||||
| `plugins` | `id` (String) | `name`, `description`, `version`, `author_id` (FK), `category`, `price_cents`, `permissions` (JSON), `status`, `s3_package_key`, `install_count`, `avg_rating` | Marketplace plugin catalog |
|
|
||||||
| `plugin_installations` | `id` (UUID) | `plugin_id` (FK), `user_id` (FK), unique constraint on (`plugin_id`, `user_id`) | Per-user install tracking |
|
|
||||||
| `plugin_reviews` | `id` (UUID) | `plugin_id` (FK), `reviewer_id` (FK), `decision`, `notes`, `reviewed_at` | Admin review decisions |
|
|
||||||
| `revenue_events` | `id` (UUID) | `plugin_id` (FK), `user_id` (FK), `amount_cents`, `developer_share_cents`, `stripe_transfer_id` | 70/30 revenue split ledger |
|
|
||||||
|
|
||||||
### Enum Types
|
|
||||||
|
|
||||||
| Enum | Values |
|
|
||||||
|---|---|
|
|
||||||
| `billing_tier` | `free`, `pro`, `power`, `team` |
|
|
||||||
| `plugin_status` | `pending_review`, `approved`, `rejected` |
|
|
||||||
| `review_decision` | `approved`, `rejected` |
|
|
||||||
|
|
||||||
### Migrations
|
|
||||||
|
|
||||||
| Version | Description |
|
|
||||||
|---|---|
|
|
||||||
| `001_initial_schema` | Creates all 9 tables with indexes and foreign key constraints |
|
|
||||||
| `002_seed_plugins` | Seeds 3 approved plugins: GitHub Sync (free), Slack Notifier (€4.99), Time Tracker (€9.99) |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## AI Agent System
|
|
||||||
|
|
||||||
The agent system uses a registry pattern with LangChain tool-calling agents powered by GPT-4o. Source: `app/agents/`, `app/core/agent_registry.py`
|
|
||||||
|
|
||||||
### Architecture
|
|
||||||
|
|
||||||
- **`BaseAgent`** — Abstract base with `user_id`, `shared_memory`, and `vector_store_context`.
|
|
||||||
- **`ChatAgent(BaseAgent)`** — Abstract `handle(query, context)` and `get_tools()` methods, plus a shared `_tool_loop(llm, messages, tools, max_iter=5)` for iterative tool calling.
|
|
||||||
- **`AgentRegistry`** — Singleton registry with `@register` decorator, `get(name)`, `list_agents()`, and `call_agent(name, query, context)`.
|
|
||||||
|
|
||||||
### Registered Agents
|
|
||||||
|
|
||||||
| Agent | Registry Name | Tools | Description |
|
|
||||||
|---|---|---|---|
|
|
||||||
| **TaskAgent** | `task_agent` | 8 | Full task and comment CRUD. Status: `todo` / `in_progress` / `done`. Priority: `high` / `medium` / `low`. Tools: `list_tasks`, `create_task`, `update_task`, `delete_task`, `list_tasks_due_today`, `list_task_comments`, `add_task_comment`, `delete_task_comment` |
|
|
||||||
| **ProjectAgent** | `project_agent` | 6 | Project lifecycle management. Status: `active` / `archived`. Prefers archiving over deletion. Tools: `list_projects`, `list_all_projects`, `get_project`, `create_project`, `update_project`, `delete_project` |
|
|
||||||
| **CheckpointAgent** | `checkpoint_agent` | 4 | Project milestones. Requires `project_id` for creation. Supports AI-suggestion and approval workflows. Tools: `list_checkpoints`, `create_checkpoint`, `update_checkpoint`, `delete_checkpoint` |
|
|
||||||
| **NoteAgent** | `note_agent` | 5 | Markdown note management. Optionally linked to projects. Tools: `list_notes`, `get_note`, `create_note`, `update_note`, `delete_note` |
|
|
||||||
|
|
||||||
All agents use the model configured by `LLM_MODEL` (default: GPT-4o) with `temperature=0` via LiteLLM. Tools return JSON action descriptors that the Electron client interprets and applies locally.
|
|
||||||
|
|
||||||
### Switching LLM Providers
|
|
||||||
|
|
||||||
The backend uses **LiteLLM** as a universal LLM gateway. All agents and the orchestrator instantiate models through a centralized factory in `app/core/llm.py`. To switch providers, change environment variables — no code changes required:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# OpenAI (default)
|
|
||||||
LLM_MODEL=gpt-4o
|
|
||||||
LLM_ROUTER_MODEL=gpt-4o-mini
|
|
||||||
|
|
||||||
# Anthropic
|
|
||||||
LLM_MODEL=anthropic/claude-3.5-sonnet
|
|
||||||
LLM_ROUTER_MODEL=anthropic/claude-3-haiku
|
|
||||||
|
|
||||||
# Google Gemini
|
|
||||||
LLM_MODEL=gemini/gemini-pro
|
|
||||||
LLM_ROUTER_MODEL=gemini/gemini-flash
|
|
||||||
|
|
||||||
# Local Ollama
|
|
||||||
LLM_MODEL=ollama/llama3
|
|
||||||
LLM_ROUTER_MODEL=ollama/llama3
|
|
||||||
|
|
||||||
# AWS Bedrock
|
|
||||||
LLM_MODEL=bedrock/anthropic.claude-v2
|
|
||||||
LLM_ROUTER_MODEL=bedrock/anthropic.claude-instant-v1
|
|
||||||
```
|
|
||||||
|
|
||||||
See the [LiteLLM provider docs](https://docs.litellm.ai/docs/providers) for the full list of 100+ supported providers and model naming conventions.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Orchestration & Execution Plans
|
|
||||||
|
|
||||||
Source: `app/core/orchestrator.py`, `app/core/execution_plan.py`
|
|
||||||
|
|
||||||
### Orchestrator
|
|
||||||
|
|
||||||
1. **`classify_intent(message, context, registry)`** — Uses the router model (`LLM_ROUTER_MODEL`, default: GPT-4o-mini) to determine which agent should handle a message. Falls back to `task_agent` when classification is ambiguous.
|
|
||||||
2. **`route_single(agent_name, message, context)`** — Routes to a single agent and returns a `ChatResponse`.
|
|
||||||
3. **`route_pipeline(agent_names, message, context)`** — Executes agents sequentially; each receives `previous_results` from earlier agents. A final LLM synthesis step merges all results.
|
|
||||||
4. **`orchestrate(request)`** — Main entry point. In `direct` mode, returns a `ChatResponse`. In `plan` mode, returns an `ExecutionPlan`.
|
|
||||||
5. **`orchestrate_stream(request)`** — Streaming variant that yields 50-character text chunks with a final JSON frame.
|
|
||||||
|
|
||||||
### Execution Plans
|
|
||||||
|
|
||||||
- **`PromptTemplateRegistry`** — Maps template IDs to server-side prompt text. Clients only ever see opaque IDs, never raw prompts.
|
|
||||||
- **`ExecutionPlanBuilder`** — Fluent builder API: `add_step()`, `add_llm_step(template_id, vars)`, `add_data_step(action, data_from_step)`. Validates step references on `build()`.
|
|
||||||
- **`PlanCache`** — LRU cache (maxsize 1000) for storing plans as reusable playbooks.
|
|
||||||
|
|
||||||
### Built-in Templates (6)
|
|
||||||
|
|
||||||
`tpl_task_agent_default`, `tpl_checkpoint_agent_default`, `tpl_project_agent_default`, `tpl_note_agent_default`, `tpl_task_extract_from_project`, `tpl_note_weekly_summary`
|
|
||||||
|
|
||||||
### Built-in Playbooks (2)
|
|
||||||
|
|
||||||
| Playbook | Description |
|
|
||||||
|---|---|
|
|
||||||
| `create_tasks_from_project` | LLM extracts actionable tasks from project context, then creates task records |
|
|
||||||
| `generate_weekly_note` | LLM generates a weekly summary, then creates a note record |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Middleware
|
|
||||||
|
|
||||||
Middleware executes in this order on each request: **TierRateLimit → Sanitizer → CORS → Router**
|
|
||||||
|
|
||||||
### JWT Authentication
|
|
||||||
|
|
||||||
Source: `app/api/middleware/auth.py`
|
|
||||||
|
|
||||||
- FastAPI dependency `get_current_user` validates the `Bearer` JWT and extracts `user_id` and `email`.
|
|
||||||
- **Live tier lookup** — The current tier is fetched from the `subscriptions` table on every request (not cached in the JWT), so upgrades and downgrades take immediate effect.
|
|
||||||
- Falls back to `free` when no subscription row exists.
|
|
||||||
- Raises `401 Unauthorized` on invalid or expired tokens.
|
|
||||||
- **Exempt paths:** `/api/v1/auth/register`, `/api/v1/auth/login`, `/api/v1/billing/webhook`
|
|
||||||
|
|
||||||
### Tier-Based Rate Limiter
|
|
||||||
|
|
||||||
Source: `app/api/middleware/rate_limit.py`
|
|
||||||
|
|
||||||
- `TierRateLimitMiddleware` — Sliding-window in-process rate limiter (no Redis dependency).
|
|
||||||
- Per-user 60-second window sized by subscription tier:
|
|
||||||
|
|
||||||
| Tier | Requests / Minute |
|
|
||||||
|---|---|
|
|
||||||
| Free | 20 |
|
|
||||||
| Pro | 60 |
|
|
||||||
| Power | 120 |
|
|
||||||
| Team | 200 |
|
|
||||||
|
|
||||||
- Returns `429 Too Many Requests` with a `Retry-After` header when the limit is exceeded.
|
|
||||||
- **Exempt paths:** register, login, webhook, health
|
|
||||||
|
|
||||||
### Response Sanitizer
|
|
||||||
|
|
||||||
Source: `app/api/middleware/sanitizer.py`
|
|
||||||
|
|
||||||
- Runs only on `/api/v1/chat` endpoints.
|
|
||||||
- Scans JSON response bodies and replaces leaked prompt IP fragments with `[REDACTED]`.
|
|
||||||
- Detects: system prompt openers, agent routing metadata, LangChain tool schemas, internal reasoning markers (`<thinking>`, `[INST]`), and known prompt fingerprints.
|
|
||||||
- Logs sanitization events as `WARNING`.
|
|
||||||
- Binary responses (storage, backup) are never touched.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Storage Layer
|
|
||||||
|
|
||||||
### Blob Store
|
|
||||||
|
|
||||||
Source: `app/storage/blob_store.py`
|
|
||||||
|
|
||||||
- S3-backed storage for E2E encrypted blobs.
|
|
||||||
- Object keys follow the pattern: `{user_id}/{table}/{record_id}`
|
|
||||||
- Server-side SSE-S3 encryption at rest (additional layer on top of client-side E2E encryption).
|
|
||||||
- Methods: `upload()`, `download()`, `delete()` (idempotent), `list_keys()`
|
|
||||||
- The backend **never inspects or decrypts blob content**.
|
|
||||||
|
|
||||||
### Vector Store
|
|
||||||
|
|
||||||
Source: `app/storage/vector_store.py`
|
|
||||||
|
|
||||||
- Runtime-configurable: **Pinecone** (when `PINECONE_API_KEY` is set) or **Qdrant** (fallback).
|
|
||||||
- User isolation: Pinecone uses `namespace=user_id`; Qdrant filters by `user_id` payload field.
|
|
||||||
- 32-dimensional SHA-256-derived float vectors (deterministic, not semantically meaningful on encrypted data — a documented trade-off for privacy).
|
|
||||||
- Encrypted blobs are stored as base64 in metadata/payload for verbatim retrieval.
|
|
||||||
- Methods: `upsert()`, `search()`, `delete()`
|
|
||||||
|
|
||||||
### Encryption Utilities
|
|
||||||
|
|
||||||
Source: `app/storage/encryption.py`
|
|
||||||
|
|
||||||
- `verify_checksum(blob, checksum)` — SHA-256 hash comparison using `hmac.compare_digest` (constant-time to prevent timing attacks).
|
|
||||||
- `reject_if_tampered(blob, checksum)` — Raises HTTP 400 on checksum mismatch.
|
|
||||||
- **No decryption key ever reaches the backend.**
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Billing & Tiers
|
|
||||||
|
|
||||||
Source: `app/billing/stripe_service.py`, `app/billing/tier_manager.py`
|
|
||||||
|
|
||||||
### Feature Matrix
|
|
||||||
|
|
||||||
| Feature | Free | Pro | Power | Team |
|
|
||||||
|---|---|---|---|---|
|
|
||||||
| AI Agents | 3 | Unlimited | Unlimited | Unlimited |
|
|
||||||
| Batch Active | 2 | 10 | Unlimited | Unlimited |
|
|
||||||
| Cloud Storage | 0 GB | 5 GB | 25 GB | Unlimited |
|
|
||||||
| Backup Storage | 0 GB | 5 GB | 25 GB | Unlimited |
|
|
||||||
| LLM Providers | 1 | Unlimited | Unlimited | Unlimited |
|
|
||||||
| Batch Builder | — | — | ✓ | ✓ |
|
|
||||||
| Plugin Marketplace | — | — | ✓ | ✓ |
|
|
||||||
| SSO | — | — | — | ✓ |
|
|
||||||
| Rate Limit | 20 req/min | 60 req/min | 120 req/min | 200 req/min |
|
|
||||||
|
|
||||||
### Stripe Integration
|
|
||||||
|
|
||||||
- **Checkout** — `create_checkout_session(user_id, tier)` creates a Stripe Checkout session. Returns a stub URL when Stripe is not configured.
|
|
||||||
- **Webhooks** — Handles `checkout.session.completed`, `customer.subscription.updated`, `customer.subscription.deleted`, and `invoice.payment_failed`.
|
|
||||||
- **Subscription management** — `get_subscription()` returns the current subscription record; `cancel_subscription()` cancels via the Stripe API and reverts the user to the free tier.
|
|
||||||
- **Price IDs:** `price_pro_monthly`, `price_power_monthly`, `price_team_monthly`
|
|
||||||
|
|
||||||
### Tier Manager
|
|
||||||
|
|
||||||
- `get_tier(user_id)` — Returns the user's current billing tier.
|
|
||||||
- `check_feature(tier, feature)` — Boolean feature gate check.
|
|
||||||
- `require_feature(tier, feature)` — Raises HTTP 403 if the feature is not available.
|
|
||||||
- `enforce_quota(user_id, tier)` / `enforce_backup_quota(user_id, tier)` — Raises HTTP 402 if storage limits are exceeded.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Plugin Marketplace
|
|
||||||
|
|
||||||
Source: `app/marketplace/`
|
|
||||||
|
|
||||||
### Plugin Registry
|
|
||||||
|
|
||||||
- PostgreSQL-backed catalog of submitted and approved plugins.
|
|
||||||
- `list_plugins(db, category, query, page, sort)` — Paginated listing (page size: 20) with optional filtering by category, text search, and sorting by `rating`, `installs`, or `newest`.
|
|
||||||
- `get_plugin(db, plugin_id)` — Full manifest with install count and ratings.
|
|
||||||
- `submit_plugin(db, manifest, s3_key)` — Submits a plugin with `pending_review` status.
|
|
||||||
- `approve_plugin()` / `reject_plugin(reason)` — Admin workflow for plugin approval.
|
|
||||||
- `record_install()` / `record_uninstall()` — Tracks per-user installations and updates install counts.
|
|
||||||
|
|
||||||
### Review Queue
|
|
||||||
|
|
||||||
- Automated security checklist before human review:
|
|
||||||
- Plugin ID must match `^[a-z0-9-]+$`
|
|
||||||
- Permissions must be from the allowed set only
|
|
||||||
- No binary blobs in the manifest
|
|
||||||
- **Allowed permissions:** `read:tasks`, `write:tasks`, `read:projects`, `write:projects`, `read:notes`, `write:notes`, `read:checkpoints`, `write:checkpoints`, `read:calendar`, `write:calendar`
|
|
||||||
- `get_pending(db)` — Lists plugins awaiting review.
|
|
||||||
- `submit_review(db, plugin_id, reviewer_id, decision, notes)` — Records the review decision.
|
|
||||||
|
|
||||||
### Revenue Sharing
|
|
||||||
|
|
||||||
- **70% developer / 30% platform** split on all paid plugin sales.
|
|
||||||
- `record_install(db, plugin_id, user_id, amount_cents)` — Records the revenue event and triggers a Stripe Connect transfer for the developer share.
|
|
||||||
- `get_earnings(db, developer_id, period)` — Aggregated earnings report for plugin developers.
|
|
||||||
- Gracefully stubs transfers when Stripe is not configured.
|
|
||||||
|
|
||||||
### Seed Plugins
|
|
||||||
|
|
||||||
| Plugin | Category | Price |
|
|
||||||
|---|---|---|
|
|
||||||
| GitHub Sync | Productivity | Free |
|
|
||||||
| Slack Notifier | Communication | €4.99 |
|
|
||||||
| Time Tracker | Productivity | €9.99 |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Testing
|
|
||||||
|
|
||||||
### Running Tests
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Run all tests
|
|
||||||
pytest
|
|
||||||
|
|
||||||
# Run a specific test file
|
|
||||||
pytest tests/test_auth.py
|
|
||||||
|
|
||||||
# Run with verbose output
|
|
||||||
pytest -v
|
|
||||||
```
|
|
||||||
|
|
||||||
### Test Infrastructure
|
|
||||||
|
|
||||||
- **Database:** Async SQLite in-memory via `aiosqlite` + `StaticPool` — fast, no PostgreSQL needed.
|
|
||||||
- **S3 mock:** `moto[s3]` with a fixture that patches `BlobStore` settings.
|
|
||||||
- **Auth helpers:** `make_jwt(tier)` and `auth_header(tier)` generate per-tier test tokens.
|
|
||||||
- **Seed data:** Auto-creates one `User` + `Subscription` per tier (free/pro/power/team) before each test.
|
|
||||||
- **Plugin seeds:** Fixture adds 3 approved plugins for marketplace tests.
|
|
||||||
- **FK enforcement:** SQLite `PRAGMA foreign_keys=ON`.
|
|
||||||
- **No external dependencies** — all tests run fully offline.
|
|
||||||
|
|
||||||
### Test Coverage
|
|
||||||
|
|
||||||
| File | Coverage |
|
|
||||||
|---|---|
|
|
||||||
| `test_auth.py` | Register, login, token access, refresh, expiration |
|
|
||||||
| `test_orchestrator.py` | Intent classification, single agent routing, pipeline, plan mode |
|
|
||||||
| `test_agents.py` | Each agent with mocked LLM: registration, tools, handle method |
|
|
||||||
| `test_storage.py` | Create, list, download, update, delete records; checksum rejection; quota enforcement |
|
|
||||||
| `test_backup.py` | Upload, download, history, delete; tier-based storage limits |
|
|
||||||
| `test_plugins.py` | List, install, uninstall, revenue events, tier gate enforcement |
|
|
||||||
| `test_agent_registry.py` | Registry singleton, registration, lookup, listing |
|
|
||||||
| `test_execution_plan.py` | Plan builder, template registry, plan cache |
|
|
||||||
| `test_middleware.py` | Rate limiting by tier, sanitizer prompt leak detection |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Project Structure
|
|
||||||
|
|
||||||
```
|
|
||||||
adiuva-api/
|
|
||||||
├── alembic.ini # Alembic configuration
|
|
||||||
├── BACKEND_PLAN.md # Architecture & design decisions
|
|
||||||
├── docker-compose.yml # Docker Compose (app + PostgreSQL)
|
|
||||||
├── Dockerfile # Multi-stage production build
|
|
||||||
├── requirements.txt # Python dependencies
|
|
||||||
│
|
|
||||||
├── alembic/ # Database migrations
|
|
||||||
│ ├── env.py # Alembic environment config
|
|
||||||
│ ├── script.py.mako # Migration template
|
|
||||||
│ └── versions/
|
|
||||||
│ ├── 001_initial_schema.py # Tables, indexes, FKs
|
|
||||||
│ └── 002_seed_plugins.py # Seed marketplace plugins
|
|
||||||
│
|
|
||||||
├── app/ # Application source
|
|
||||||
│ ├── main.py # FastAPI app factory, middleware, routes
|
|
||||||
│ ├── db.py # Async SQLAlchemy engine & session
|
|
||||||
│ ├── models.py # SQLAlchemy ORM models (9 tables)
|
|
||||||
│ ├── schemas.py # Pydantic request/response schemas
|
|
||||||
│ │
|
|
||||||
│ ├── config/
|
|
||||||
│ │ └── settings.py # Pydantic Settings (env vars)
|
|
||||||
│ │
|
|
||||||
│ ├── agents/ # LLM-powered domain agents
|
|
||||||
│ │ ├── task_agent.py # Task & comment CRUD (8 tools)
|
|
||||||
│ │ ├── project_agent.py # Project lifecycle (6 tools)
|
|
||||||
│ │ ├── checkpoint_agent.py # Milestones (4 tools)
|
|
||||||
│ │ └── note_agent.py # Markdown notes (5 tools)
|
|
||||||
│ │
|
|
||||||
│ ├── core/ # Orchestration engine
|
|
||||||
│ │ ├── agent_registry.py # BaseAgent, ChatAgent, AgentRegistry
|
|
||||||
│ │ ├── llm.py # LiteLLM factory (get_llm, get_router_llm)
|
|
||||||
│ │ ├── orchestrator.py # Intent classification & routing
|
|
||||||
│ │ └── execution_plan.py # Plan builder, templates, cache
|
|
||||||
│ │
|
|
||||||
│ ├── api/ # HTTP layer
|
|
||||||
│ │ ├── deps.py # Shared FastAPI dependencies
|
|
||||||
│ │ ├── middleware/
|
|
||||||
│ │ │ ├── auth.py # JWT validation, live tier lookup
|
|
||||||
│ │ │ ├── rate_limit.py # Sliding-window tier rate limiter
|
|
||||||
│ │ │ └── sanitizer.py # Prompt IP leak protection
|
|
||||||
│ │ └── routes/
|
|
||||||
│ │ ├── auth.py # Register, login, refresh, me
|
|
||||||
│ │ ├── chat.py # Chat + WebSocket streaming
|
|
||||||
│ │ ├── plans.py # Execution plan playbooks
|
|
||||||
│ │ ├── storage.py # E2E encrypted record CRUD
|
|
||||||
│ │ ├── vectors.py # Vector upsert, search, delete
|
|
||||||
│ │ ├── backup.py # Encrypted backup management
|
|
||||||
│ │ ├── plugins.py # Marketplace browse & install
|
|
||||||
│ │ └── billing.py # Stripe checkout & webhooks
|
|
||||||
│ │
|
|
||||||
│ ├── storage/ # Storage backends
|
|
||||||
│ │ ├── blob_store.py # S3 blob storage
|
|
||||||
│ │ ├── vector_store.py # Pinecone / Qdrant vector store
|
|
||||||
│ │ └── encryption.py # Checksum verification utilities
|
|
||||||
│ │
|
|
||||||
│ ├── billing/ # Subscription management
|
|
||||||
│ │ ├── stripe_service.py # Stripe API integration
|
|
||||||
│ │ └── tier_manager.py # Feature matrix & quota enforcement
|
|
||||||
│ │
|
|
||||||
│ └── marketplace/ # Plugin ecosystem
|
|
||||||
│ ├── plugin_registry.py # Catalog CRUD & search
|
|
||||||
│ ├── plugin_review.py # Security checklist & review queue
|
|
||||||
│ └── revenue_share.py # 70/30 split & Stripe Connect
|
|
||||||
│
|
|
||||||
└── tests/ # Test suite
|
|
||||||
├── conftest.py # Fixtures: DB, S3, auth, seeds
|
|
||||||
├── test_auth.py
|
|
||||||
├── test_orchestrator.py
|
|
||||||
├── test_agents.py
|
|
||||||
├── test_storage.py
|
|
||||||
├── test_backup.py
|
|
||||||
├── test_plugins.py
|
|
||||||
├── test_agent_registry.py
|
|
||||||
├── test_execution_plan.py
|
|
||||||
└── test_middleware.py
|
|
||||||
```
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## License
|
|
||||||
|
|
||||||
*To be determined.*
|
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ import re
|
|||||||
from logging.config import fileConfig
|
from logging.config import fileConfig
|
||||||
|
|
||||||
from alembic import context
|
from alembic import context
|
||||||
from sqlalchemy import engine_from_config, pool
|
from sqlalchemy import pool
|
||||||
from sqlalchemy.ext.asyncio import create_async_engine
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||||||
|
|
||||||
# Alembic Config object (gives access to alembic.ini values).
|
# Alembic Config object (gives access to alembic.ini values).
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
"""Initial schema: users, refresh_tokens, subscriptions, storage_records,
|
"""Initial schema: users, refresh_tokens, subscriptions.
|
||||||
backup_metadata, plugins, plugin_installations, plugin_reviews, revenue_events.
|
|
||||||
|
|
||||||
Revision ID: 001
|
Revision ID: 001
|
||||||
Revises:
|
Revises:
|
||||||
@@ -21,18 +20,13 @@ depends_on: Union[str, Sequence[str], None] = None
|
|||||||
|
|
||||||
|
|
||||||
def upgrade() -> None:
|
def upgrade() -> None:
|
||||||
# ── Enum types ────────────────────────────────────────────────────────
|
# ── Enum types — idempotent creation via exception handling ───────────
|
||||||
billing_tier = postgresql.ENUM(
|
op.execute("""
|
||||||
"free", "pro", "power", "team", name="billing_tier", create_type=False
|
DO $$ BEGIN
|
||||||
)
|
CREATE TYPE billing_tier AS ENUM ('free', 'pro', 'power', 'team');
|
||||||
plugin_status = postgresql.ENUM(
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
"pending_review", "approved", "rejected", name="plugin_status", create_type=False
|
END $$;
|
||||||
)
|
""")
|
||||||
review_decision = postgresql.ENUM(
|
|
||||||
"approved", "rejected", name="review_decision", create_type=False
|
|
||||||
)
|
|
||||||
for enum in (billing_tier, plugin_status, review_decision):
|
|
||||||
enum.create(op.get_bind(), checkfirst=True)
|
|
||||||
|
|
||||||
# ── users ─────────────────────────────────────────────────────────────
|
# ── users ─────────────────────────────────────────────────────────────
|
||||||
op.create_table(
|
op.create_table(
|
||||||
@@ -40,7 +34,7 @@ def upgrade() -> None:
|
|||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
sa.Column("email", sa.String(255), nullable=False),
|
sa.Column("email", sa.String(255), nullable=False),
|
||||||
sa.Column("password_hash", sa.String(255), nullable=False),
|
sa.Column("password_hash", sa.String(255), nullable=False),
|
||||||
sa.Column("tier", sa.Enum("free", "pro", "power", "team", name="billing_tier", create_type=False), nullable=False, server_default="free"),
|
sa.Column("tier", postgresql.ENUM("free", "pro", "power", "team", name="billing_tier", create_type=False), nullable=False, server_default="free"),
|
||||||
sa.Column("stripe_customer_id", sa.String(255), nullable=True),
|
sa.Column("stripe_customer_id", sa.String(255), nullable=True),
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
@@ -70,7 +64,7 @@ def upgrade() -> None:
|
|||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
sa.Column("stripe_subscription_id", sa.String(255), nullable=True),
|
sa.Column("stripe_subscription_id", sa.String(255), nullable=True),
|
||||||
sa.Column("tier", sa.Enum("free", "pro", "power", "team", name="billing_tier", create_type=False), nullable=False, server_default="free"),
|
sa.Column("tier", postgresql.ENUM("free", "pro", "power", "team", name="billing_tier", create_type=False), nullable=False, server_default="free"),
|
||||||
sa.Column("status", sa.String(50), nullable=False, server_default="free"),
|
sa.Column("status", sa.String(50), nullable=False, server_default="free"),
|
||||||
sa.Column("current_period_end", sa.DateTime(timezone=True), nullable=True),
|
sa.Column("current_period_end", sa.DateTime(timezone=True), nullable=True),
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
@@ -81,122 +75,10 @@ def upgrade() -> None:
|
|||||||
op.create_index("ix_subscriptions_user_id", "subscriptions", ["user_id"])
|
op.create_index("ix_subscriptions_user_id", "subscriptions", ["user_id"])
|
||||||
op.create_index("ix_subscriptions_stripe_id", "subscriptions", ["stripe_subscription_id"])
|
op.create_index("ix_subscriptions_stripe_id", "subscriptions", ["stripe_subscription_id"])
|
||||||
|
|
||||||
# ── storage_records ───────────────────────────────────────────────────
|
|
||||||
op.create_table(
|
|
||||||
"storage_records",
|
|
||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("table_name", sa.String(100), nullable=False),
|
|
||||||
sa.Column("s3_key", sa.String(500), nullable=False),
|
|
||||||
sa.Column("checksum", sa.String(64), nullable=False),
|
|
||||||
sa.Column("size_bytes", sa.Integer, nullable=False),
|
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.PrimaryKeyConstraint("id"),
|
|
||||||
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
|
||||||
)
|
|
||||||
op.create_index("ix_storage_records_user_id", "storage_records", ["user_id"])
|
|
||||||
|
|
||||||
# ── backup_metadata ───────────────────────────────────────────────────
|
|
||||||
op.create_table(
|
|
||||||
"backup_metadata",
|
|
||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("s3_key", sa.String(500), nullable=False),
|
|
||||||
sa.Column("version", sa.Integer, nullable=False),
|
|
||||||
sa.Column("timestamp", sa.BigInteger, nullable=False),
|
|
||||||
sa.Column("checksum", sa.String(64), nullable=False),
|
|
||||||
sa.Column("size_bytes", sa.Integer, nullable=False),
|
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.PrimaryKeyConstraint("id"),
|
|
||||||
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
|
||||||
)
|
|
||||||
op.create_index("ix_backup_metadata_user_id", "backup_metadata", ["user_id"])
|
|
||||||
|
|
||||||
# ── plugins ───────────────────────────────────────────────────────────
|
|
||||||
op.create_table(
|
|
||||||
"plugins",
|
|
||||||
sa.Column("id", sa.String(255), nullable=False),
|
|
||||||
sa.Column("name", sa.String(255), nullable=False),
|
|
||||||
sa.Column("description", sa.Text, nullable=False, server_default=""),
|
|
||||||
sa.Column("version", sa.String(50), nullable=False, server_default="1.0.0"),
|
|
||||||
sa.Column("author_id", postgresql.UUID(as_uuid=False), nullable=True),
|
|
||||||
sa.Column("author_name", sa.String(255), nullable=False, server_default=""),
|
|
||||||
sa.Column("category", sa.String(100), nullable=False, server_default=""),
|
|
||||||
sa.Column("price_cents", sa.Integer, nullable=False, server_default="0"),
|
|
||||||
sa.Column("permissions", sa.Text, nullable=False, server_default="[]"),
|
|
||||||
sa.Column("status", sa.Enum("pending_review", "approved", "rejected", name="plugin_status", create_type=False), nullable=False, server_default="pending_review"),
|
|
||||||
sa.Column("s3_package_key", sa.String(500), nullable=True),
|
|
||||||
sa.Column("install_count", sa.Integer, nullable=False, server_default="0"),
|
|
||||||
sa.Column("avg_rating", sa.Float, nullable=False, server_default="0.0"),
|
|
||||||
sa.Column("rejection_reason", sa.Text, nullable=True),
|
|
||||||
sa.Column("submitted_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.PrimaryKeyConstraint("id"),
|
|
||||||
sa.ForeignKeyConstraint(["author_id"], ["users.id"], ondelete="SET NULL"),
|
|
||||||
)
|
|
||||||
|
|
||||||
# ── plugin_installations ──────────────────────────────────────────────
|
|
||||||
op.create_table(
|
|
||||||
"plugin_installations",
|
|
||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("plugin_id", sa.String(255), nullable=False),
|
|
||||||
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("installed_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.PrimaryKeyConstraint("id"),
|
|
||||||
sa.ForeignKeyConstraint(["plugin_id"], ["plugins.id"], ondelete="CASCADE"),
|
|
||||||
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
|
||||||
sa.UniqueConstraint("plugin_id", "user_id", name="uq_plugin_user"),
|
|
||||||
)
|
|
||||||
op.create_index("ix_plugin_installations_plugin_id", "plugin_installations", ["plugin_id"])
|
|
||||||
op.create_index("ix_plugin_installations_user_id", "plugin_installations", ["user_id"])
|
|
||||||
|
|
||||||
# ── plugin_reviews ────────────────────────────────────────────────────
|
|
||||||
op.create_table(
|
|
||||||
"plugin_reviews",
|
|
||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("plugin_id", sa.String(255), nullable=False),
|
|
||||||
sa.Column("reviewer_id", postgresql.UUID(as_uuid=False), nullable=True),
|
|
||||||
sa.Column("decision", sa.Enum("approved", "rejected", name="review_decision", create_type=False), nullable=False),
|
|
||||||
sa.Column("notes", sa.Text, nullable=True),
|
|
||||||
sa.Column("reviewed_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.PrimaryKeyConstraint("id"),
|
|
||||||
sa.ForeignKeyConstraint(["plugin_id"], ["plugins.id"], ondelete="CASCADE"),
|
|
||||||
sa.ForeignKeyConstraint(["reviewer_id"], ["users.id"], ondelete="SET NULL"),
|
|
||||||
)
|
|
||||||
op.create_index("ix_plugin_reviews_plugin_id", "plugin_reviews", ["plugin_id"])
|
|
||||||
|
|
||||||
# ── revenue_events ────────────────────────────────────────────────────
|
|
||||||
op.create_table(
|
|
||||||
"revenue_events",
|
|
||||||
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("plugin_id", sa.String(255), nullable=False),
|
|
||||||
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
|
||||||
sa.Column("amount_cents", sa.Integer, nullable=False, server_default="0"),
|
|
||||||
sa.Column("developer_share_cents", sa.Integer, nullable=False, server_default="0"),
|
|
||||||
sa.Column("stripe_transfer_id", sa.String(255), nullable=True),
|
|
||||||
sa.Column("paid_at", sa.DateTime(timezone=True), nullable=True),
|
|
||||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
|
||||||
sa.PrimaryKeyConstraint("id"),
|
|
||||||
sa.ForeignKeyConstraint(["plugin_id"], ["plugins.id"], ondelete="CASCADE"),
|
|
||||||
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
|
||||||
)
|
|
||||||
op.create_index("ix_revenue_events_plugin_id", "revenue_events", ["plugin_id"])
|
|
||||||
op.create_index("ix_revenue_events_user_id", "revenue_events", ["user_id"])
|
|
||||||
|
|
||||||
|
|
||||||
def downgrade() -> None:
|
def downgrade() -> None:
|
||||||
op.drop_table("revenue_events")
|
|
||||||
op.drop_table("plugin_reviews")
|
|
||||||
op.drop_table("plugin_installations")
|
|
||||||
op.drop_table("plugins")
|
|
||||||
op.drop_table("backup_metadata")
|
|
||||||
op.drop_table("storage_records")
|
|
||||||
op.drop_table("subscriptions")
|
op.drop_table("subscriptions")
|
||||||
op.drop_table("refresh_tokens")
|
op.drop_table("refresh_tokens")
|
||||||
op.drop_table("users")
|
op.drop_table("users")
|
||||||
|
|
||||||
op.execute("DROP TYPE IF EXISTS review_decision")
|
|
||||||
op.execute("DROP TYPE IF EXISTS plugin_status")
|
|
||||||
op.execute("DROP TYPE IF EXISTS billing_tier")
|
op.execute("DROP TYPE IF EXISTS billing_tier")
|
||||||
|
|||||||
@@ -1,92 +0,0 @@
|
|||||||
"""Seed approved plugins: GitHub Sync, Slack Notifier, Time Tracker.
|
|
||||||
|
|
||||||
Revision ID: 002
|
|
||||||
Revises: 001
|
|
||||||
Create Date: 2026-03-03
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import json
|
|
||||||
from datetime import datetime, timezone
|
|
||||||
from typing import Sequence, Union
|
|
||||||
|
|
||||||
import sqlalchemy as sa
|
|
||||||
from alembic import op
|
|
||||||
|
|
||||||
revision: str = "002"
|
|
||||||
down_revision: Union[str, None] = "001"
|
|
||||||
branch_labels: Union[str, Sequence[str], None] = None
|
|
||||||
depends_on: Union[str, Sequence[str], None] = None
|
|
||||||
|
|
||||||
_SEED_PLUGINS = [
|
|
||||||
{
|
|
||||||
"id": "plugin-github-sync",
|
|
||||||
"name": "GitHub Sync",
|
|
||||||
"description": "Sync tasks with GitHub Issues and pull requests.",
|
|
||||||
"version": "1.0.0",
|
|
||||||
"author_name": "Adiuva",
|
|
||||||
"category": "productivity",
|
|
||||||
"price_cents": 0,
|
|
||||||
"permissions": json.dumps(["read:tasks", "write:tasks"]),
|
|
||||||
"status": "approved",
|
|
||||||
"s3_package_key": "plugins/plugin-github-sync/1.0.0/package.zip",
|
|
||||||
"install_count": 0,
|
|
||||||
"avg_rating": 0.0,
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"id": "plugin-slack-notify",
|
|
||||||
"name": "Slack Notifier",
|
|
||||||
"description": "Post task and checkpoint updates to Slack channels.",
|
|
||||||
"version": "1.2.0",
|
|
||||||
"author_name": "Adiuva",
|
|
||||||
"category": "communication",
|
|
||||||
"price_cents": 499,
|
|
||||||
"permissions": json.dumps(["read:tasks", "read:checkpoints"]),
|
|
||||||
"status": "approved",
|
|
||||||
"s3_package_key": "plugins/plugin-slack-notify/1.2.0/package.zip",
|
|
||||||
"install_count": 0,
|
|
||||||
"avg_rating": 0.0,
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"id": "plugin-time-tracker",
|
|
||||||
"name": "Time Tracker",
|
|
||||||
"description": "Track time spent on tasks with automatic reporting.",
|
|
||||||
"version": "0.9.1",
|
|
||||||
"author_name": "Third Party",
|
|
||||||
"category": "productivity",
|
|
||||||
"price_cents": 999,
|
|
||||||
"permissions": json.dumps(["read:tasks", "write:tasks"]),
|
|
||||||
"status": "approved",
|
|
||||||
"s3_package_key": "plugins/plugin-time-tracker/0.9.1/package.zip",
|
|
||||||
"install_count": 0,
|
|
||||||
"avg_rating": 0.0,
|
|
||||||
},
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
def upgrade() -> None:
|
|
||||||
plugins = sa.table(
|
|
||||||
"plugins",
|
|
||||||
sa.column("id", sa.String),
|
|
||||||
sa.column("name", sa.String),
|
|
||||||
sa.column("description", sa.Text),
|
|
||||||
sa.column("version", sa.String),
|
|
||||||
sa.column("author_name", sa.String),
|
|
||||||
sa.column("category", sa.String),
|
|
||||||
sa.column("price_cents", sa.Integer),
|
|
||||||
sa.column("permissions", sa.Text),
|
|
||||||
sa.column("status", sa.Enum("pending_review", "approved", "rejected", name="plugin_status")),
|
|
||||||
sa.column("s3_package_key", sa.String),
|
|
||||||
sa.column("install_count", sa.Integer),
|
|
||||||
sa.column("avg_rating", sa.Float),
|
|
||||||
)
|
|
||||||
op.bulk_insert(plugins, _SEED_PLUGINS)
|
|
||||||
|
|
||||||
|
|
||||||
def downgrade() -> None:
|
|
||||||
op.execute(
|
|
||||||
"DELETE FROM plugins WHERE id IN ("
|
|
||||||
"'plugin-github-sync', 'plugin-slack-notify', 'plugin-time-tracker'"
|
|
||||||
")"
|
|
||||||
)
|
|
||||||
127
alembic/versions/003_agent_tables.py
Normal file
127
alembic/versions/003_agent_tables.py
Normal file
@@ -0,0 +1,127 @@
|
|||||||
|
"""Add agent config and run log tables: local_agent_configs, cloud_agent_configs, agent_run_logs.
|
||||||
|
|
||||||
|
Revision ID: 003
|
||||||
|
Revises: 002
|
||||||
|
Create Date: 2026-03-05
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
revision: str = "003"
|
||||||
|
down_revision: Union[str, None] = "001"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
# ── Enum types — idempotent creation ──────────────────────────────────
|
||||||
|
op.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE agent_type AS ENUM ('local', 'cloud');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
op.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE agent_run_status AS ENUM ('running', 'success', 'error', 'partial');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
op.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE cloud_provider AS ENUM ('gmail', 'teams', 'outlook');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
|
||||||
|
# ── local_agent_configs ───────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"local_agent_configs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("device_id", sa.String(255), nullable=False),
|
||||||
|
sa.Column("name", sa.String(255), nullable=False),
|
||||||
|
sa.Column("directory_paths", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("data_types", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("prompt_template", sa.Text, nullable=False, server_default=""),
|
||||||
|
sa.Column("file_extensions", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("schedule_cron", sa.String(100), nullable=False, server_default="0 */6 * * *"),
|
||||||
|
sa.Column("enabled", sa.Boolean, nullable=False, server_default=sa.true()),
|
||||||
|
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_local_agent_configs_user_id", "local_agent_configs", ["user_id"])
|
||||||
|
|
||||||
|
# ── cloud_agent_configs ───────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"cloud_agent_configs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"provider",
|
||||||
|
postgresql.ENUM("gmail", "teams", "outlook", name="cloud_provider", create_type=False),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("name", sa.String(255), nullable=False),
|
||||||
|
sa.Column("data_types", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("prompt_template", sa.Text, nullable=False, server_default=""),
|
||||||
|
sa.Column("oauth_token_encrypted", sa.Text, nullable=True),
|
||||||
|
sa.Column("filter_config", sa.JSON, nullable=True),
|
||||||
|
sa.Column("schedule_cron", sa.String(100), nullable=False, server_default="0 */6 * * *"),
|
||||||
|
sa.Column("enabled", sa.Boolean, nullable=False, server_default=sa.true()),
|
||||||
|
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_cloud_agent_configs_user_id", "cloud_agent_configs", ["user_id"])
|
||||||
|
|
||||||
|
# ── agent_run_logs ─────────────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"agent_run_logs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
# Plain string — not a FK because it references either local_agent_configs or
|
||||||
|
# cloud_agent_configs depending on agent_type.
|
||||||
|
sa.Column("agent_id", sa.String(255), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"agent_type",
|
||||||
|
postgresql.ENUM("local", "cloud", name="agent_type", create_type=False),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"status",
|
||||||
|
postgresql.ENUM("running", "success", "error", "partial", name="agent_run_status", create_type=False),
|
||||||
|
nullable=False,
|
||||||
|
server_default="running",
|
||||||
|
),
|
||||||
|
sa.Column("items_processed", sa.Integer, nullable=False, server_default="0"),
|
||||||
|
sa.Column("items_created", sa.Integer, nullable=False, server_default="0"),
|
||||||
|
sa.Column("errors", sa.JSON, nullable=True),
|
||||||
|
sa.Column("started_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_agent_run_logs_user_id", "agent_run_logs", ["user_id"])
|
||||||
|
op.create_index("ix_agent_run_logs_agent_id", "agent_run_logs", ["agent_id"])
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_table("agent_run_logs")
|
||||||
|
op.drop_table("cloud_agent_configs")
|
||||||
|
op.drop_table("local_agent_configs")
|
||||||
|
|
||||||
|
op.execute("DROP TYPE IF EXISTS cloud_provider;")
|
||||||
|
op.execute("DROP TYPE IF EXISTS agent_run_status;")
|
||||||
|
op.execute("DROP TYPE IF EXISTS agent_type;")
|
||||||
144
alembic/versions/004_add_memory_tables.py
Normal file
144
alembic/versions/004_add_memory_tables.py
Normal file
@@ -0,0 +1,144 @@
|
|||||||
|
"""Add memory tables and user encryption_key column.
|
||||||
|
|
||||||
|
Memory tables:
|
||||||
|
memory_core — per-user key/value preferences (encrypted)
|
||||||
|
memory_associative — semantic memory with pgvector embedding (encrypted)
|
||||||
|
memory_episodic — session summaries (encrypted)
|
||||||
|
memory_proactive — behavioral patterns (encrypted)
|
||||||
|
|
||||||
|
Also adds encryption_key column to users table.
|
||||||
|
|
||||||
|
Revision ID: 004
|
||||||
|
Revises: 003
|
||||||
|
Create Date: 2026-03-08
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
revision: str = "004"
|
||||||
|
down_revision: Union[str, None] = "003"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
# ── Enable pgvector extension (idempotent) ────────────────────────────────
|
||||||
|
op.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
||||||
|
|
||||||
|
# ── Add encryption_key to users ───────────────────────────────────────────
|
||||||
|
op.add_column(
|
||||||
|
"users",
|
||||||
|
sa.Column("encryption_key", sa.String(64), nullable=True),
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── memory_core ───────────────────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"memory_core",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
||||||
|
sa.Column(
|
||||||
|
"user_id",
|
||||||
|
postgresql.UUID(as_uuid=False),
|
||||||
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("key", sa.String(255), nullable=False),
|
||||||
|
sa.Column("value_encrypted", sa.Text, nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"updated_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.func.now(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
op.create_index("ix_memory_core_user_id", "memory_core", ["user_id"])
|
||||||
|
|
||||||
|
# ── memory_associative ────────────────────────────────────────────────────
|
||||||
|
# The embedding column uses pgvector's vector(1536) type.
|
||||||
|
op.create_table(
|
||||||
|
"memory_associative",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
||||||
|
sa.Column(
|
||||||
|
"user_id",
|
||||||
|
postgresql.UUID(as_uuid=False),
|
||||||
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("content_encrypted", sa.Text, nullable=False),
|
||||||
|
sa.Column("entity_type", sa.String(100), nullable=True),
|
||||||
|
sa.Column("entity_id", sa.String(255), nullable=True),
|
||||||
|
sa.Column(
|
||||||
|
"updated_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.func.now(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
# Add the pgvector column separately (not supported by generic sa types)
|
||||||
|
op.execute(
|
||||||
|
"ALTER TABLE memory_associative ADD COLUMN embedding vector(1536);"
|
||||||
|
)
|
||||||
|
op.create_index("ix_memory_associative_user_id", "memory_associative", ["user_id"])
|
||||||
|
# IVFFlat index for approximate nearest-neighbour search
|
||||||
|
op.execute(
|
||||||
|
"CREATE INDEX ix_memory_associative_embedding "
|
||||||
|
"ON memory_associative USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);"
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── memory_episodic ───────────────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"memory_episodic",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
||||||
|
sa.Column(
|
||||||
|
"user_id",
|
||||||
|
postgresql.UUID(as_uuid=False),
|
||||||
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("summary_encrypted", sa.Text, nullable=False),
|
||||||
|
sa.Column("session_id", sa.String(255), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"created_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.func.now(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
op.create_index("ix_memory_episodic_user_id", "memory_episodic", ["user_id"])
|
||||||
|
op.create_index("ix_memory_episodic_session_id", "memory_episodic", ["session_id"])
|
||||||
|
|
||||||
|
# ── memory_proactive ──────────────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"memory_proactive",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
||||||
|
sa.Column(
|
||||||
|
"user_id",
|
||||||
|
postgresql.UUID(as_uuid=False),
|
||||||
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("pattern_encrypted", sa.Text, nullable=False),
|
||||||
|
sa.Column("confidence", sa.Float, nullable=False, server_default="0.5"),
|
||||||
|
sa.Column("source", sa.String(50), nullable=False, server_default="inferred"),
|
||||||
|
sa.Column(
|
||||||
|
"created_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.func.now(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
op.create_index("ix_memory_proactive_user_id", "memory_proactive", ["user_id"])
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_table("memory_proactive")
|
||||||
|
op.drop_table("memory_episodic")
|
||||||
|
op.drop_index("ix_memory_associative_embedding", "memory_associative")
|
||||||
|
op.drop_table("memory_associative")
|
||||||
|
op.drop_table("memory_core")
|
||||||
|
op.drop_column("users", "encryption_key")
|
||||||
54
alembic/versions/005_associative_pgvector.py
Normal file
54
alembic/versions/005_associative_pgvector.py
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
"""Phase 1 — confirm pgvector activation on memory_associative.
|
||||||
|
|
||||||
|
Migration 004 created the embedding column as vector(1536) and added the
|
||||||
|
IVFFlat index. This migration is the Phase-1 checkpoint:
|
||||||
|
1. Ensures the pgvector extension is enabled (idempotent).
|
||||||
|
2. Ensures the canonical Phase-1 IVFFlat index exists under the name
|
||||||
|
memory_associative_embedding_idx (creates it only if absent).
|
||||||
|
|
||||||
|
Revision ID: 005
|
||||||
|
Revises: 9a1f2d0b6c7e
|
||||||
|
Create Date: 2026-04-15
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
revision: str = "005"
|
||||||
|
down_revision: Union[str, None] = "e04100e88ace"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
# Ensure pgvector extension is enabled (also done in 004, idempotent).
|
||||||
|
op.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
||||||
|
|
||||||
|
# Ensure the canonical Phase-1 IVFFlat index exists.
|
||||||
|
# 004 may have created ix_memory_associative_embedding; this adds the
|
||||||
|
# Phase-1 name memory_associative_embedding_idx if it is missing.
|
||||||
|
op.execute(
|
||||||
|
"""
|
||||||
|
DO $$
|
||||||
|
BEGIN
|
||||||
|
IF NOT EXISTS (
|
||||||
|
SELECT 1
|
||||||
|
FROM pg_indexes
|
||||||
|
WHERE tablename = 'memory_associative'
|
||||||
|
AND indexname = 'memory_associative_embedding_idx'
|
||||||
|
) THEN
|
||||||
|
CREATE INDEX memory_associative_embedding_idx
|
||||||
|
ON memory_associative
|
||||||
|
USING ivfflat (embedding vector_cosine_ops)
|
||||||
|
WITH (lists = 100);
|
||||||
|
END IF;
|
||||||
|
END $$;
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.execute("DROP INDEX IF EXISTS memory_associative_embedding_idx;")
|
||||||
74
alembic/versions/006_memory_relations.py
Normal file
74
alembic/versions/006_memory_relations.py
Normal file
@@ -0,0 +1,74 @@
|
|||||||
|
"""Add memory_relations table (Phase 3 — relational tier).
|
||||||
|
|
||||||
|
Revision ID: 006
|
||||||
|
Revises: 1f5975a4f3f4
|
||||||
|
Create Date: 2026-04-16
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
revision: str = "006"
|
||||||
|
down_revision: Union[str, None] = "1f5975a4f3f4"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.create_table(
|
||||||
|
"memory_relations",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
||||||
|
sa.Column(
|
||||||
|
"user_id",
|
||||||
|
postgresql.UUID(as_uuid=False),
|
||||||
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("subject_label", sa.String(128), nullable=False),
|
||||||
|
sa.Column("subject_type", sa.String(32), nullable=False),
|
||||||
|
sa.Column("predicate", sa.String(64), nullable=False),
|
||||||
|
sa.Column("object_label", sa.String(128), nullable=False),
|
||||||
|
sa.Column("object_type", sa.String(32), nullable=False),
|
||||||
|
sa.Column("confidence", sa.Float, nullable=False, server_default="0.7"),
|
||||||
|
sa.Column(
|
||||||
|
"source_episode_id",
|
||||||
|
postgresql.UUID(as_uuid=False),
|
||||||
|
sa.ForeignKey("memory_episodic.id", ondelete="SET NULL"),
|
||||||
|
nullable=True,
|
||||||
|
),
|
||||||
|
sa.Column("notes_encrypted", sa.LargeBinary, nullable=True),
|
||||||
|
sa.Column(
|
||||||
|
"created_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.func.now(),
|
||||||
|
),
|
||||||
|
sa.Column(
|
||||||
|
"updated_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.func.now(),
|
||||||
|
),
|
||||||
|
sa.Column("last_confirmed_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
)
|
||||||
|
op.create_index(
|
||||||
|
"memory_relations_user_subject_idx",
|
||||||
|
"memory_relations",
|
||||||
|
["user_id", "subject_label"],
|
||||||
|
)
|
||||||
|
op.create_index(
|
||||||
|
"memory_relations_user_predicate_idx",
|
||||||
|
"memory_relations",
|
||||||
|
["user_id", "predicate"],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_index("memory_relations_user_predicate_idx", "memory_relations")
|
||||||
|
op.drop_index("memory_relations_user_subject_idx", "memory_relations")
|
||||||
|
op.drop_table("memory_relations")
|
||||||
38
alembic/versions/1f5975a4f3f4_add_extraction_queue.py
Normal file
38
alembic/versions/1f5975a4f3f4_add_extraction_queue.py
Normal file
@@ -0,0 +1,38 @@
|
|||||||
|
"""add extraction_queue
|
||||||
|
|
||||||
|
Revision ID: 1f5975a4f3f4
|
||||||
|
Revises: 005
|
||||||
|
Create Date: 2026-04-16 17:26:25.790870
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = '1f5975a4f3f4'
|
||||||
|
down_revision: Union[str, None] = '005'
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.create_table(
|
||||||
|
'extraction_queue',
|
||||||
|
sa.Column('id', sa.Uuid(as_uuid=False), nullable=False),
|
||||||
|
sa.Column('user_id', sa.Uuid(as_uuid=False), nullable=False),
|
||||||
|
sa.Column('episode_id', sa.Uuid(as_uuid=False), nullable=True),
|
||||||
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
|
||||||
|
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE'),
|
||||||
|
sa.PrimaryKeyConstraint('id'),
|
||||||
|
)
|
||||||
|
op.create_index(op.f('ix_extraction_queue_user_id'), 'extraction_queue', ['user_id'], unique=False)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_index(op.f('ix_extraction_queue_user_id'), table_name='extraction_queue')
|
||||||
|
op.drop_table('extraction_queue')
|
||||||
@@ -0,0 +1,30 @@
|
|||||||
|
"""add name and surname to users table
|
||||||
|
|
||||||
|
Revision ID: 818478c251dc
|
||||||
|
Revises: 004
|
||||||
|
Create Date: 2026-03-10 15:10:42.811947
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = '818478c251dc'
|
||||||
|
down_revision: Union[str, None] = '004'
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.add_column('users', sa.Column('name', sa.String(length=100), nullable=True))
|
||||||
|
op.add_column('users', sa.Column('surname', sa.String(length=100), nullable=True))
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_column('users', 'surname')
|
||||||
|
op.drop_column('users', 'name')
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
"""Deprecate backend agent config tables.
|
||||||
|
|
||||||
|
The Electron client is now the source of truth for agent configuration
|
||||||
|
(directory, extract targets, batch interval, custom prompt). Backend keeps
|
||||||
|
billing checks and trigger/run logs only.
|
||||||
|
|
||||||
|
Revision ID: 9a1f2d0b6c7e
|
||||||
|
Revises: 818478c251dc
|
||||||
|
Create Date: 2026-03-16
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
revision: str = "9a1f2d0b6c7e"
|
||||||
|
down_revision: Union[str, None] = "818478c251dc"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
bind = op.get_bind()
|
||||||
|
inspector = sa.inspect(bind)
|
||||||
|
existing = set(inspector.get_table_names())
|
||||||
|
|
||||||
|
if "cloud_agent_configs" in existing:
|
||||||
|
op.drop_index("ix_cloud_agent_configs_user_id", table_name="cloud_agent_configs")
|
||||||
|
op.drop_table("cloud_agent_configs")
|
||||||
|
|
||||||
|
if "local_agent_configs" in existing:
|
||||||
|
op.drop_index("ix_local_agent_configs_user_id", table_name="local_agent_configs")
|
||||||
|
op.drop_table("local_agent_configs")
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.create_table(
|
||||||
|
"local_agent_configs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("device_id", sa.String(255), nullable=False),
|
||||||
|
sa.Column("name", sa.String(255), nullable=False),
|
||||||
|
sa.Column("directory_paths", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("data_types", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("prompt_template", sa.Text, nullable=False, server_default=""),
|
||||||
|
sa.Column("file_extensions", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("schedule_cron", sa.String(100), nullable=False, server_default="0 */6 * * *"),
|
||||||
|
sa.Column("enabled", sa.Boolean, nullable=False, server_default=sa.true()),
|
||||||
|
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_local_agent_configs_user_id", "local_agent_configs", ["user_id"])
|
||||||
|
|
||||||
|
op.execute(
|
||||||
|
"""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE cloud_provider AS ENUM ('gmail', 'teams', 'outlook');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
|
op.create_table(
|
||||||
|
"cloud_agent_configs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"provider",
|
||||||
|
postgresql.ENUM("gmail", "teams", "outlook", name="cloud_provider", create_type=False),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("name", sa.String(255), nullable=False),
|
||||||
|
sa.Column("data_types", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("prompt_template", sa.Text, nullable=False, server_default=""),
|
||||||
|
sa.Column("oauth_token_encrypted", sa.Text, nullable=True),
|
||||||
|
sa.Column("filter_config", sa.JSON, nullable=True),
|
||||||
|
sa.Column("schedule_cron", sa.String(100), nullable=False, server_default="0 */6 * * *"),
|
||||||
|
sa.Column("enabled", sa.Boolean, nullable=False, server_default=sa.true()),
|
||||||
|
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_cloud_agent_configs_user_id", "cloud_agent_configs", ["user_id"])
|
||||||
@@ -0,0 +1,107 @@
|
|||||||
|
"""Restore agent config tables and add agent_config column.
|
||||||
|
|
||||||
|
9a1f2d0b6c7e dropped local_agent_configs and cloud_agent_configs, but both
|
||||||
|
ORM models are still active. This migration recreates them with agent_config
|
||||||
|
added to local_agent_configs.
|
||||||
|
|
||||||
|
Revision ID: a3b9c0d1e2f3
|
||||||
|
Revises: 9a1f2d0b6c7e
|
||||||
|
Create Date: 2026-04-07 00:00:00.000000
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = "a3b9c0d1e2f3"
|
||||||
|
down_revision: Union[str, None] = "9a1f2d0b6c7e"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
# Recreate enum types (idempotent — they may already exist from migration 003)
|
||||||
|
op.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE agent_type AS ENUM ('local', 'cloud');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
op.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE agent_run_status AS ENUM ('running', 'success', 'error', 'partial');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
op.execute("""
|
||||||
|
DO $$ BEGIN
|
||||||
|
CREATE TYPE cloud_provider AS ENUM ('gmail', 'teams', 'outlook');
|
||||||
|
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||||
|
END $$;
|
||||||
|
""")
|
||||||
|
|
||||||
|
bind = op.get_bind()
|
||||||
|
inspector = sa.inspect(bind)
|
||||||
|
existing = set(inspector.get_table_names())
|
||||||
|
|
||||||
|
# ── local_agent_configs (with agent_config column) ────────────────────
|
||||||
|
if "local_agent_configs" not in existing:
|
||||||
|
op.create_table(
|
||||||
|
"local_agent_configs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("device_id", sa.String(255), nullable=False),
|
||||||
|
sa.Column("name", sa.String(255), nullable=False),
|
||||||
|
sa.Column("directory_paths", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("data_types", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("prompt_template", sa.Text, nullable=False, server_default=""),
|
||||||
|
sa.Column("agent_config", sa.JSON, nullable=True),
|
||||||
|
sa.Column("file_extensions", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("schedule_cron", sa.String(100), nullable=False, server_default="0 */6 * * *"),
|
||||||
|
sa.Column("enabled", sa.Boolean, nullable=False, server_default=sa.true()),
|
||||||
|
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_local_agent_configs_user_id", "local_agent_configs", ["user_id"])
|
||||||
|
|
||||||
|
# ── cloud_agent_configs ───────────────────────────────────────────────
|
||||||
|
if "cloud_agent_configs" not in existing:
|
||||||
|
op.create_table(
|
||||||
|
"cloud_agent_configs",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"provider",
|
||||||
|
postgresql.ENUM("gmail", "teams", "outlook", name="cloud_provider", create_type=False),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column("name", sa.String(255), nullable=False),
|
||||||
|
sa.Column("data_types", sa.JSON, nullable=False, server_default="[]"),
|
||||||
|
sa.Column("prompt_template", sa.Text, nullable=False, server_default=""),
|
||||||
|
sa.Column("oauth_token_encrypted", sa.Text, nullable=True),
|
||||||
|
sa.Column("filter_config", sa.JSON, nullable=True),
|
||||||
|
sa.Column("schedule_cron", sa.String(100), nullable=False, server_default="0 */6 * * *"),
|
||||||
|
sa.Column("enabled", sa.Boolean, nullable=False, server_default=sa.true()),
|
||||||
|
sa.Column("last_run_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_cloud_agent_configs_user_id", "cloud_agent_configs", ["user_id"])
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_index("ix_cloud_agent_configs_user_id", table_name="cloud_agent_configs")
|
||||||
|
op.drop_table("cloud_agent_configs")
|
||||||
|
op.drop_index("ix_local_agent_configs_user_id", table_name="local_agent_configs")
|
||||||
|
op.drop_table("local_agent_configs")
|
||||||
56
alembic/versions/b4c0d1e2f3a4_add_oauth_and_avatar.py
Normal file
56
alembic/versions/b4c0d1e2f3a4_add_oauth_and_avatar.py
Normal file
@@ -0,0 +1,56 @@
|
|||||||
|
"""Add oauth_accounts table, nullable password_hash, avatar_url to users.
|
||||||
|
|
||||||
|
Revision ID: b4c0d1e2f3a4
|
||||||
|
Revises: a3b9c0d1e2f3
|
||||||
|
Create Date: 2026-04-10 00:00:00.000000
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = "b4c0d1e2f3a4"
|
||||||
|
down_revision: Union[str, None] = "a3b9c0d1e2f3"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
# ── users: make password_hash nullable (social users have no password) ──
|
||||||
|
op.alter_column("users", "password_hash", existing_type=sa.String(255), nullable=True)
|
||||||
|
|
||||||
|
# ── users: add avatar_url ─────────────────────────────────────────────
|
||||||
|
op.add_column("users", sa.Column("avatar_url", sa.String(2048), nullable=True))
|
||||||
|
|
||||||
|
# ── oauth_accounts ────────────────────────────────────────────────────
|
||||||
|
op.create_table(
|
||||||
|
"oauth_accounts",
|
||||||
|
sa.Column("id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("user_id", postgresql.UUID(as_uuid=False), nullable=False),
|
||||||
|
sa.Column("provider", sa.String(50), nullable=False),
|
||||||
|
sa.Column("provider_user_id", sa.String(255), nullable=False),
|
||||||
|
sa.Column("provider_email", sa.String(255), nullable=True),
|
||||||
|
sa.Column(
|
||||||
|
"created_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
nullable=False,
|
||||||
|
server_default=sa.text("now()"),
|
||||||
|
),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||||
|
sa.UniqueConstraint("provider", "provider_user_id", name="uq_oauth_provider_user"),
|
||||||
|
)
|
||||||
|
op.create_index("ix_oauth_accounts_user_id", "oauth_accounts", ["user_id"])
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_index("ix_oauth_accounts_user_id", table_name="oauth_accounts")
|
||||||
|
op.drop_table("oauth_accounts")
|
||||||
|
op.drop_column("users", "avatar_url")
|
||||||
|
op.alter_column("users", "password_hash", existing_type=sa.String(255), nullable=False)
|
||||||
31
alembic/versions/c5d1e2f3a4b5_add_onboarding_completed_at.py
Normal file
31
alembic/versions/c5d1e2f3a4b5_add_onboarding_completed_at.py
Normal file
@@ -0,0 +1,31 @@
|
|||||||
|
"""Add onboarding_completed_at column to users table.
|
||||||
|
|
||||||
|
Revision ID: c5d1e2f3a4b5
|
||||||
|
Revises: b4c0d1e2f3a4
|
||||||
|
Create Date: 2026-04-11 00:00:00.000000
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = "c5d1e2f3a4b5"
|
||||||
|
down_revision: Union[str, None] = "b4c0d1e2f3a4"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.add_column(
|
||||||
|
"users",
|
||||||
|
sa.Column("onboarding_completed_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_column("users", "onboarding_completed_at")
|
||||||
46
alembic/versions/d6e3f4a5b6c7_folder_index_tables.py
Normal file
46
alembic/versions/d6e3f4a5b6c7_folder_index_tables.py
Normal file
@@ -0,0 +1,46 @@
|
|||||||
|
"""Add token tracking columns for folder integration.
|
||||||
|
|
||||||
|
Revision ID: d6e3f4a5b6c7
|
||||||
|
Revises: 006
|
||||||
|
Create Date: 2026-05-11 00:00:00.000000
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
from sqlalchemy.dialects.postgresql import UUID
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = "d6e3f4a5b6c7"
|
||||||
|
down_revision: Union[str, None] = "006"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.add_column(
|
||||||
|
"agent_run_logs",
|
||||||
|
sa.Column("tokens_used", sa.Integer(), nullable=False, server_default="0"),
|
||||||
|
)
|
||||||
|
op.create_table(
|
||||||
|
"monthly_token_usage",
|
||||||
|
sa.Column("user_id", UUID(as_uuid=False), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
|
||||||
|
sa.Column("year_month", sa.String(7), nullable=False),
|
||||||
|
sa.Column("feature", sa.String(64), nullable=False),
|
||||||
|
sa.Column("tokens_used", sa.Integer(), nullable=False, server_default="0"),
|
||||||
|
sa.PrimaryKeyConstraint("user_id", "year_month", "feature"),
|
||||||
|
)
|
||||||
|
op.create_index(
|
||||||
|
"ix_monthly_token_usage_user_month",
|
||||||
|
"monthly_token_usage",
|
||||||
|
["user_id", "year_month"],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_index("ix_monthly_token_usage_user_month", table_name="monthly_token_usage")
|
||||||
|
op.drop_table("monthly_token_usage")
|
||||||
|
op.drop_column("agent_run_logs", "tokens_used")
|
||||||
34
alembic/versions/e04100e88ace_avatar_url_varchar_to_text.py
Normal file
34
alembic/versions/e04100e88ace_avatar_url_varchar_to_text.py
Normal file
@@ -0,0 +1,34 @@
|
|||||||
|
"""avatar_url_varchar_to_text
|
||||||
|
|
||||||
|
Revision ID: e04100e88ace
|
||||||
|
Revises: c5d1e2f3a4b5
|
||||||
|
Create Date: 2026-04-13 09:13:06.733674
|
||||||
|
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = 'e04100e88ace'
|
||||||
|
down_revision: Union[str, None] = 'c5d1e2f3a4b5'
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.alter_column('users', 'avatar_url',
|
||||||
|
existing_type=sa.VARCHAR(length=2048),
|
||||||
|
type_=sa.Text(),
|
||||||
|
existing_nullable=True)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.alter_column('users', 'avatar_url',
|
||||||
|
existing_type=sa.Text(),
|
||||||
|
type_=sa.VARCHAR(length=2048),
|
||||||
|
existing_nullable=True)
|
||||||
@@ -1,5 +1,5 @@
|
|||||||
"""Import all agent modules to trigger @registry.register decorators."""
|
"""Expose tool modules used by deep orchestrator-worker graphs."""
|
||||||
|
|
||||||
from app.agents import checkpoint_agent, note_agent, project_agent, task_agent
|
from app.agents import filesystem_agent, timeline_agent, note_agent, project_agent, task_agent
|
||||||
|
|
||||||
__all__ = ["checkpoint_agent", "note_agent", "project_agent", "task_agent"]
|
__all__ = ["filesystem_agent", "timeline_agent", "note_agent", "project_agent", "task_agent"]
|
||||||
|
|||||||
@@ -1,127 +0,0 @@
|
|||||||
"""Checkpoint agent — project milestone management (list, create, update, delete)."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
from langchain_core.messages import HumanMessage, SystemMessage
|
|
||||||
from langchain_core.tools import tool
|
|
||||||
|
|
||||||
from app.core.agent_registry import ChatAgent, registry
|
|
||||||
from app.core.llm import get_llm
|
|
||||||
from app.core.ws_context import execute_on_client
|
|
||||||
|
|
||||||
_SYSTEM_PROMPT = (
|
|
||||||
"You are a project checkpoint assistant. Checkpoints are milestone dates that\n"
|
|
||||||
"track progress on a project — they are not calendar events.\n\n"
|
|
||||||
"Rules:\n"
|
|
||||||
" - project_id is REQUIRED for every create; confirm with the user if unknown\n"
|
|
||||||
" - date is a Unix timestamp in milliseconds; convert human-readable dates\n"
|
|
||||||
" - is_ai_suggested: 1 when proactively proposing a checkpoint, 0 otherwise\n"
|
|
||||||
" - is_approved: 0 until the user explicitly confirms; then 1\n"
|
|
||||||
" - For update_checkpoint, use -1 for integer fields you do not want to change\n"
|
|
||||||
" - Listing without a project_id returns all checkpoints across projects\n"
|
|
||||||
" - Always echo the title and formatted date in your confirmation."
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@tool
|
|
||||||
async def list_checkpoints(project_id: str = "") -> str:
|
|
||||||
"""List checkpoints. Provide project_id to scope to a specific project."""
|
|
||||||
result = await execute_on_client(
|
|
||||||
action="select",
|
|
||||||
table="checkpoints",
|
|
||||||
filters={"projectId": project_id or None},
|
|
||||||
)
|
|
||||||
rows = result.get("rows", [])
|
|
||||||
if not rows:
|
|
||||||
return "No checkpoints found."
|
|
||||||
lines = [f"- {r['title']} (date: {r['date']}, id: {r['id']})" for r in rows]
|
|
||||||
return f"Found {len(rows)} checkpoint(s):\n" + "\n".join(lines)
|
|
||||||
|
|
||||||
|
|
||||||
@tool
|
|
||||||
async def create_checkpoint(
|
|
||||||
project_id: str,
|
|
||||||
title: str,
|
|
||||||
date: int,
|
|
||||||
is_ai_suggested: int = 0,
|
|
||||||
is_approved: int = 0,
|
|
||||||
) -> str:
|
|
||||||
"""Create a project checkpoint (milestone).
|
|
||||||
project_id: REQUIRED UUID of the parent project
|
|
||||||
title: descriptive name for the milestone
|
|
||||||
date: Unix timestamp in milliseconds
|
|
||||||
is_ai_suggested: 1 if proactively suggested, 0 if user-requested
|
|
||||||
is_approved: 0 until the user confirms
|
|
||||||
"""
|
|
||||||
result = await execute_on_client(
|
|
||||||
action="insert",
|
|
||||||
table="checkpoints",
|
|
||||||
data={
|
|
||||||
"projectId": project_id,
|
|
||||||
"title": title,
|
|
||||||
"date": date,
|
|
||||||
"isAiSuggested": is_ai_suggested,
|
|
||||||
"isApproved": is_approved,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
row = result["row"]
|
|
||||||
return f"Checkpoint created: '{row['title']}' (id: {row['id']}, date: {row['date']})"
|
|
||||||
|
|
||||||
|
|
||||||
@tool
|
|
||||||
async def update_checkpoint(
|
|
||||||
checkpoint_id: str,
|
|
||||||
title: str = "",
|
|
||||||
date: int = -1,
|
|
||||||
is_approved: int = -1,
|
|
||||||
) -> str:
|
|
||||||
"""Update a checkpoint. Only pass fields that should change.
|
|
||||||
checkpoint_id: UUID of the checkpoint (required)
|
|
||||||
date: -1 means unchanged; any other value sets the new date (ms timestamp)
|
|
||||||
is_approved: -1 means unchanged; 0 or 1 sets the approval state
|
|
||||||
"""
|
|
||||||
updates: dict[str, Any] = {}
|
|
||||||
if title:
|
|
||||||
updates["title"] = title
|
|
||||||
if date != -1:
|
|
||||||
updates["date"] = date
|
|
||||||
if is_approved != -1:
|
|
||||||
updates["isApproved"] = is_approved
|
|
||||||
result = await execute_on_client(
|
|
||||||
action="update",
|
|
||||||
table="checkpoints",
|
|
||||||
data={"id": checkpoint_id, "updates": updates},
|
|
||||||
)
|
|
||||||
row = result["row"]
|
|
||||||
return f"Checkpoint updated: '{row['title']}' (id: {row['id']})"
|
|
||||||
|
|
||||||
|
|
||||||
@tool
|
|
||||||
async def delete_checkpoint(checkpoint_id: str) -> str:
|
|
||||||
"""Delete a checkpoint permanently by its UUID."""
|
|
||||||
await execute_on_client(action="delete", table="checkpoints", data={"id": checkpoint_id})
|
|
||||||
return f"Checkpoint {checkpoint_id} deleted."
|
|
||||||
|
|
||||||
|
|
||||||
@registry.register
|
|
||||||
class CheckpointAgent(ChatAgent):
|
|
||||||
def get_name(self) -> str:
|
|
||||||
return "checkpoint_agent"
|
|
||||||
|
|
||||||
def get_description(self) -> str:
|
|
||||||
return "Manages project checkpoints (milestones): list, create, update, delete"
|
|
||||||
|
|
||||||
def get_tools(self) -> list[Any]:
|
|
||||||
return [list_checkpoints, create_checkpoint, update_checkpoint, delete_checkpoint]
|
|
||||||
|
|
||||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
|
||||||
llm = get_llm()
|
|
||||||
messages = [
|
|
||||||
SystemMessage(content=_SYSTEM_PROMPT),
|
|
||||||
HumanMessage(
|
|
||||||
content=f"User query: {query}\nContext: {json.dumps(context)[:1000]}"
|
|
||||||
),
|
|
||||||
]
|
|
||||||
return await self._tool_loop(llm, messages, self.get_tools())
|
|
||||||
52
app/agents/client_agent.py
Normal file
52
app/agents/client_agent.py
Normal file
@@ -0,0 +1,52 @@
|
|||||||
|
"""Client agent — read-only tools for the clients table."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def list_clients(search: str = "", limit: int = 20) -> str:
|
||||||
|
"""List clients, optionally filtered by a name/email substring search.
|
||||||
|
|
||||||
|
search: optional substring to match against client name or email.
|
||||||
|
limit: max rows to return (default 20).
|
||||||
|
"""
|
||||||
|
filters: dict[str, Any] = {"limit": limit}
|
||||||
|
if search:
|
||||||
|
filters["search"] = search
|
||||||
|
|
||||||
|
result = await execute_on_client(action="select", table="clients", filters=filters)
|
||||||
|
rows = result.get("rows", [])
|
||||||
|
if not rows:
|
||||||
|
return "No clients found."
|
||||||
|
lines = [
|
||||||
|
f"- {r.get('name', '?')} (id: {r.get('id')}, email: {r.get('email', '')}, "
|
||||||
|
f"company: {r.get('company', '')})"
|
||||||
|
for r in rows
|
||||||
|
]
|
||||||
|
return f"Found {len(rows)} client(s):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def get_client(id: str) -> str:
|
||||||
|
"""Get full details for one client by UUID.
|
||||||
|
|
||||||
|
id: the client's UUID.
|
||||||
|
"""
|
||||||
|
if not id:
|
||||||
|
return "Client id is required."
|
||||||
|
|
||||||
|
result = await execute_on_client(action="get", table="clients", data={"id": id})
|
||||||
|
row = result.get("row") or result.get("rows", [None])[0] if result else None
|
||||||
|
if not row:
|
||||||
|
return f"Client '{id}' not found."
|
||||||
|
return f"Client details:\n{json.dumps(row, ensure_ascii=False, indent=2)}"
|
||||||
|
|
||||||
|
|
||||||
|
CLIENT_TOOLS: list[Any] = [list_clients, get_client]
|
||||||
194
app/agents/filesystem_agent.py
Normal file
194
app/agents/filesystem_agent.py
Normal file
@@ -0,0 +1,194 @@
|
|||||||
|
"""Filesystem agent — tools for reading local directories and files on Electron.
|
||||||
|
|
||||||
|
These tools delegate to the Electron client via ``execute_on_client()`` using
|
||||||
|
the same WS tool-call round-trip pattern as CRUD tools. The Electron app
|
||||||
|
handles actual disk I/O and responds with ``tool_result`` frames.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
|
# Max characters returned by read_file_content in journey (exploration) tools.
|
||||||
|
# The journey only needs to understand file structure, not full content.
|
||||||
|
_JOURNEY_READ_MAX_CHARS: int = 4000
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_path(path: str, base: str) -> str:
|
||||||
|
"""Resolve *path* against *base* when *path* is relative.
|
||||||
|
|
||||||
|
The LLM often passes ``"."`` meaning "the configured directory".
|
||||||
|
Without this, Electron resolves ``"."`` relative to its own CWD instead
|
||||||
|
of the user's chosen directory.
|
||||||
|
"""
|
||||||
|
if os.path.isabs(path):
|
||||||
|
return path
|
||||||
|
return str(Path(base) / path)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def list_directory(path: str) -> str:
|
||||||
|
"""List files and folders in a local directory on the user's device.
|
||||||
|
|
||||||
|
Returns a formatted listing of entries with name, type (file/directory),
|
||||||
|
and full path.
|
||||||
|
"""
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="list_directory",
|
||||||
|
data={"path": path},
|
||||||
|
)
|
||||||
|
entries: list[dict[str, Any]] = result.get("entries", [])
|
||||||
|
if not entries:
|
||||||
|
return f"Directory '{path}' is empty or does not exist."
|
||||||
|
lines: list[str] = []
|
||||||
|
for entry in entries:
|
||||||
|
entry_type = entry.get("type", "unknown")
|
||||||
|
entry_name = entry.get("name", "")
|
||||||
|
entry_path = entry.get("path", "")
|
||||||
|
lines.append(f"- [{entry_type}] {entry_name} ({entry_path})")
|
||||||
|
return f"Directory listing for '{path}' ({len(entries)} entries):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def read_file_content(path: str) -> str:
|
||||||
|
"""Read the text content of a local file on the user's device.
|
||||||
|
|
||||||
|
Returns the file content as a string. Large files may be truncated
|
||||||
|
by the Electron client.
|
||||||
|
"""
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="read_file_content",
|
||||||
|
data={"path": path},
|
||||||
|
)
|
||||||
|
content: str = result.get("content", "")
|
||||||
|
if not content:
|
||||||
|
return f"File '{path}' is empty or could not be read."
|
||||||
|
return content
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def get_file_metadata(path: str) -> str:
|
||||||
|
"""Get metadata for a local file: size, creation date, modification date, extension.
|
||||||
|
|
||||||
|
Returns a formatted summary of the file's metadata.
|
||||||
|
"""
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="get_file_metadata",
|
||||||
|
data={"path": path},
|
||||||
|
)
|
||||||
|
size = result.get("size", "unknown")
|
||||||
|
created = result.get("createdAt", "unknown")
|
||||||
|
modified = result.get("modifiedAt", "unknown")
|
||||||
|
extension = result.get("extension", "unknown")
|
||||||
|
name = result.get("name", path)
|
||||||
|
return (
|
||||||
|
f"File: {name}\n"
|
||||||
|
f" Extension: {extension}\n"
|
||||||
|
f" Size: {size} bytes\n"
|
||||||
|
f" Created: {created}\n"
|
||||||
|
f" Modified: {modified}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
FILESYSTEM_TOOLS: list[Any] = [
|
||||||
|
list_directory,
|
||||||
|
read_file_content,
|
||||||
|
get_file_metadata,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def make_directory_tools(base_directory: str) -> list[Any]:
|
||||||
|
"""Return filesystem tools that resolve relative paths against *base_directory*.
|
||||||
|
|
||||||
|
Use this instead of ``FILESYSTEM_TOOLS`` whenever you know the user's target
|
||||||
|
directory upfront (e.g., journey setup sessions). Relative paths like ``"."``
|
||||||
|
from the LLM are resolved to the correct absolute path before being sent to
|
||||||
|
the Electron client, preventing it from falling back to its own CWD.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def _compact_for_journey(raw: str) -> str:
|
||||||
|
"""Strip HTML noise and truncate for journey exploration.
|
||||||
|
|
||||||
|
The journey LLM only needs to understand file structure (headers,
|
||||||
|
first paragraphs). Full CSS/style blocks are pure noise that eat
|
||||||
|
up context window budget.
|
||||||
|
"""
|
||||||
|
text = re.sub(r"<style[^>]*>.*?</style>", "", raw, flags=re.DOTALL | re.IGNORECASE)
|
||||||
|
text = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL | re.IGNORECASE)
|
||||||
|
text = re.sub(r"<!--.*?-->", "", text, flags=re.DOTALL)
|
||||||
|
if len(text) > _JOURNEY_READ_MAX_CHARS:
|
||||||
|
text = text[:_JOURNEY_READ_MAX_CHARS] + "\n[…truncated for exploration]"
|
||||||
|
return text
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def list_directory(path: str) -> str: # noqa: F811
|
||||||
|
"""List files and folders in a local directory on the user's device.
|
||||||
|
|
||||||
|
Returns a formatted listing of entries with name, type (file/directory),
|
||||||
|
and full path.
|
||||||
|
"""
|
||||||
|
resolved = _resolve_path(path, base_directory)
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="list_directory",
|
||||||
|
data={"path": resolved},
|
||||||
|
)
|
||||||
|
entries: list[dict[str, Any]] = result.get("entries", [])
|
||||||
|
if not entries:
|
||||||
|
return f"Directory '{resolved}' is empty or does not exist."
|
||||||
|
lines: list[str] = []
|
||||||
|
for entry in entries:
|
||||||
|
entry_type = entry.get("type", "unknown")
|
||||||
|
entry_name = entry.get("name", "")
|
||||||
|
entry_path = entry.get("path", "")
|
||||||
|
lines.append(f"- [{entry_type}] {entry_name} ({entry_path})")
|
||||||
|
return f"Directory listing for '{resolved}' ({len(entries)} entries):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def read_file_content(path: str) -> str: # noqa: F811
|
||||||
|
"""Read the text content of a local file on the user's device.
|
||||||
|
|
||||||
|
Returns the file content as a string. Large files may be truncated
|
||||||
|
by the Electron client.
|
||||||
|
"""
|
||||||
|
resolved = _resolve_path(path, base_directory)
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="read_file_content",
|
||||||
|
data={"path": resolved},
|
||||||
|
)
|
||||||
|
content: str = result.get("content", "")
|
||||||
|
if not content:
|
||||||
|
return f"File '{resolved}' is empty or could not be read."
|
||||||
|
return _compact_for_journey(content)
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def get_file_metadata(path: str) -> str: # noqa: F811
|
||||||
|
"""Get metadata for a local file: size, creation date, modification date, extension.
|
||||||
|
|
||||||
|
Returns a formatted summary of the file's metadata.
|
||||||
|
"""
|
||||||
|
resolved = _resolve_path(path, base_directory)
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="get_file_metadata",
|
||||||
|
data={"path": resolved},
|
||||||
|
)
|
||||||
|
size = result.get("size", "unknown")
|
||||||
|
created = result.get("createdAt", "unknown")
|
||||||
|
modified = result.get("modifiedAt", "unknown")
|
||||||
|
extension = result.get("extension", "unknown")
|
||||||
|
name = result.get("name", resolved)
|
||||||
|
return (
|
||||||
|
f"File: {name}\n"
|
||||||
|
f" Extension: {extension}\n"
|
||||||
|
f" Size: {size} bytes\n"
|
||||||
|
f" Created: {created}\n"
|
||||||
|
f" Modified: {modified}"
|
||||||
|
)
|
||||||
|
|
||||||
|
return [list_directory, read_file_content, get_file_metadata]
|
||||||
168
app/agents/folder_agent.py
Normal file
168
app/agents/folder_agent.py
Normal file
@@ -0,0 +1,168 @@
|
|||||||
|
"""Scoped file-read and search tools for the project folder feature."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
|
from app.core.folder_indexer import _extract_docx_text, _extract_pdf_text
|
||||||
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
|
# Cap returned slice size to keep tool output under control.
|
||||||
|
_MAX_RETURN_CHARS = 50_000
|
||||||
|
_MAX_SEARCH_MATCHES = 20
|
||||||
|
|
||||||
|
|
||||||
|
def _is_unsafe_path(rel: str) -> bool:
|
||||||
|
if not rel:
|
||||||
|
return True
|
||||||
|
norm = rel.replace("\\", "/")
|
||||||
|
if norm.startswith("/"):
|
||||||
|
return True
|
||||||
|
# Windows drive letter
|
||||||
|
if len(rel) >= 2 and rel[1] == ":":
|
||||||
|
return True
|
||||||
|
parts = norm.split("/")
|
||||||
|
return ".." in parts
|
||||||
|
|
||||||
|
|
||||||
|
async def _fetch_file(project_id: str, relative_path: str, offset: int, length: int) -> dict:
|
||||||
|
"""Return the raw Electron tool_result dict for a file read."""
|
||||||
|
return await execute_on_client(
|
||||||
|
action="read_project_folder_file",
|
||||||
|
data={
|
||||||
|
"projectId": project_id,
|
||||||
|
"relativePath": relative_path,
|
||||||
|
"offset": offset,
|
||||||
|
"length": length,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _decode(result: dict) -> tuple[str, str, int]:
|
||||||
|
"""Decode a tool_result into (text, kind, total_size). For pdf/docx,
|
||||||
|
extracts text from base64. For images, returns a placeholder string.
|
||||||
|
For text, content is already a sliced utf-8 string.
|
||||||
|
"""
|
||||||
|
kind = result.get("kind", "text")
|
||||||
|
content = result.get("content", "") or ""
|
||||||
|
total = int(result.get("totalSize", 0) or 0)
|
||||||
|
if kind == "image":
|
||||||
|
return ("[Image file — cannot be navigated as text. See manifest summary.]", kind, total)
|
||||||
|
if kind == "pdf":
|
||||||
|
return (_extract_pdf_text(content), kind, total)
|
||||||
|
if kind == "docx":
|
||||||
|
return (_extract_docx_text(content), kind, total)
|
||||||
|
return (content, kind, total)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def read_project_folder_file(
|
||||||
|
project_id: str,
|
||||||
|
relative_path: str,
|
||||||
|
offset: int = 0,
|
||||||
|
length: int = _MAX_RETURN_CHARS,
|
||||||
|
) -> str:
|
||||||
|
"""Read a slice of a file inside the project's linked folder.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
project_id: project ID.
|
||||||
|
relative_path: path relative to the linked folder root.
|
||||||
|
offset: char offset to start reading from (0 = beginning).
|
||||||
|
length: max chars to return. Default 50000. Use smaller values to save tokens.
|
||||||
|
|
||||||
|
Returns text content slice with a header showing position. Header tells you
|
||||||
|
when more content is available; call again with the suggested next offset.
|
||||||
|
|
||||||
|
For PDF / DOCX files the backend extracts text first, then applies offset/length
|
||||||
|
on the extracted text. For images returns a placeholder; navigate with the
|
||||||
|
manifest summary instead.
|
||||||
|
"""
|
||||||
|
if _is_unsafe_path(relative_path):
|
||||||
|
return "Access denied"
|
||||||
|
|
||||||
|
result = await _fetch_file(project_id, relative_path, offset, length)
|
||||||
|
text, kind, total_size = _decode(result)
|
||||||
|
|
||||||
|
if not text and kind in ("missing", "error"):
|
||||||
|
return f"File not found or unreadable: {relative_path}"
|
||||||
|
|
||||||
|
if kind in ("pdf", "docx"):
|
||||||
|
# Backend extracted full text — apply offset/length on chars.
|
||||||
|
sliced = text[offset:offset + length]
|
||||||
|
slice_end = min(offset + length, len(text))
|
||||||
|
header = (
|
||||||
|
f"[file={relative_path} kind={kind} offset={offset} end={slice_end} "
|
||||||
|
f"totalChars={len(text)}]"
|
||||||
|
)
|
||||||
|
if slice_end < len(text):
|
||||||
|
header += f"\n[More content available — call again with offset={slice_end}.]"
|
||||||
|
return header + "\n" + sliced
|
||||||
|
|
||||||
|
if kind == "text":
|
||||||
|
slice_end = offset + len(text)
|
||||||
|
header = (
|
||||||
|
f"[file={relative_path} kind=text offset={offset} end={slice_end} "
|
||||||
|
f"totalBytes={total_size}]"
|
||||||
|
)
|
||||||
|
if slice_end < total_size:
|
||||||
|
header += f"\n[More content available — call again with offset={slice_end}.]"
|
||||||
|
return header + "\n" + text
|
||||||
|
|
||||||
|
# image or unknown
|
||||||
|
return text
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def search_project_folder_file(
|
||||||
|
project_id: str,
|
||||||
|
relative_path: str,
|
||||||
|
query: str,
|
||||||
|
context_lines: int = 3,
|
||||||
|
) -> str:
|
||||||
|
"""Search a project folder file for a query string (case-insensitive substring).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
project_id: project ID.
|
||||||
|
relative_path: path relative to the linked folder root.
|
||||||
|
query: text to search for.
|
||||||
|
context_lines: number of lines of context around each match (default 3).
|
||||||
|
|
||||||
|
Returns matching line ranges with surrounding context and 1-based line numbers.
|
||||||
|
Capped at 20 matches; if more exist the header shows the total.
|
||||||
|
|
||||||
|
Works on text, code, markdown, PDF (extracted), and DOCX (extracted).
|
||||||
|
Images and binary files are not searchable.
|
||||||
|
"""
|
||||||
|
if _is_unsafe_path(relative_path):
|
||||||
|
return "Access denied"
|
||||||
|
if not query:
|
||||||
|
return "Empty query."
|
||||||
|
|
||||||
|
# For text we still need full file; pass length=very large.
|
||||||
|
result = await _fetch_file(project_id, relative_path, offset=0, length=10_000_000)
|
||||||
|
text, kind, _ = _decode(result)
|
||||||
|
|
||||||
|
if not text and kind in ("missing", "error"):
|
||||||
|
return f"File not found or unreadable: {relative_path}"
|
||||||
|
if kind == "image":
|
||||||
|
return "Cannot search inside images."
|
||||||
|
|
||||||
|
lines = text.splitlines()
|
||||||
|
q = query.lower()
|
||||||
|
matches = [i for i, line in enumerate(lines) if q in line.lower()]
|
||||||
|
if not matches:
|
||||||
|
return f"No matches for '{query}' in {relative_path}."
|
||||||
|
|
||||||
|
shown = matches[:_MAX_SEARCH_MATCHES]
|
||||||
|
snippets: list[str] = []
|
||||||
|
for i in shown:
|
||||||
|
start = max(0, i - context_lines)
|
||||||
|
end = min(len(lines), i + context_lines + 1)
|
||||||
|
block = "\n".join(f"{n + 1:5d}: {lines[n]}" for n in range(start, end))
|
||||||
|
snippets.append(block)
|
||||||
|
|
||||||
|
header = f"[file={relative_path} matches={len(matches)} showing={len(shown)} query='{query}']"
|
||||||
|
body = "\n---\n".join(snippets)
|
||||||
|
return header + "\n" + body
|
||||||
|
|
||||||
|
|
||||||
|
FOLDER_TOOLS = [read_project_folder_file, search_project_folder_file]
|
||||||
@@ -1,43 +1,50 @@
|
|||||||
"""Note agent — Markdown note management (list, get, create, update, delete)."""
|
"""Note agent — Markdown note management (list, get, create, update, propose edit)."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import re
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langchain_core.messages import HumanMessage, SystemMessage
|
|
||||||
from langchain_core.tools import tool
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
from app.core.agent_registry import ChatAgent, registry
|
from app.core.note_summarizer import generate_note_summary
|
||||||
from app.core.llm import embed, get_llm
|
|
||||||
from app.core.ws_context import execute_on_client
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
_SYSTEM_PROMPT = (
|
_UUID_RE = re.compile(
|
||||||
"You are a note-taking assistant. You help users create, retrieve, update,\n"
|
r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$"
|
||||||
"and delete Markdown notes in their workspace.\n\n"
|
|
||||||
"Rules:\n"
|
|
||||||
" - content is always Markdown; preserve formatting when updating\n"
|
|
||||||
" - project_id is optional; link a note to a project when mentioned\n"
|
|
||||||
" - When updating, call get_note first if you need to read existing content\n"
|
|
||||||
" before appending or replacing sections\n"
|
|
||||||
" - list_notes without project_id returns all notes; scope with project_id\n"
|
|
||||||
" when the user is working within a specific project\n"
|
|
||||||
" - Do not fabricate note content — reflect what the user provides or what\n"
|
|
||||||
" is already in the note (retrieved via get_note)."
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_uuid(value: str) -> bool:
|
||||||
|
return bool(_UUID_RE.match(value))
|
||||||
|
|
||||||
|
|
||||||
|
def _fmt_summary(row: dict) -> str:
|
||||||
|
summary = (row.get("aiSummary") or row.get("ai_summary") or "").strip()
|
||||||
|
if summary:
|
||||||
|
return f" — {summary}"
|
||||||
|
snippet = (row.get("content") or "")[:120].replace("\n", " ").strip()
|
||||||
|
return f" — {snippet}" if snippet else ""
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
async def list_notes(project_id: str = "") -> str:
|
async def list_notes(project_id: str = "") -> str:
|
||||||
"""List notes, optionally scoped to a project by project_id."""
|
"""List notes with AI summaries, optionally scoped to a project by project_id.
|
||||||
|
|
||||||
|
Returns id, title, and ai_summary for each note so you can decide which
|
||||||
|
note to read in full with get_note before creating or updating.
|
||||||
|
"""
|
||||||
|
normalized_project_id = project_id if (project_id and _is_uuid(project_id)) else ""
|
||||||
result = await execute_on_client(
|
result = await execute_on_client(
|
||||||
action="select",
|
action="select",
|
||||||
table="notes",
|
table="notes",
|
||||||
filters={"projectId": project_id or None},
|
filters={"projectId": normalized_project_id or None},
|
||||||
)
|
)
|
||||||
rows = result.get("rows", [])
|
rows = result.get("rows", [])
|
||||||
if not rows:
|
if not rows:
|
||||||
return "No notes found."
|
return "No notes found."
|
||||||
lines = [f"- {r['title']} (id: {r['id']})" for r in rows]
|
lines = [f" - [{r['id']}] {r['title']}{_fmt_summary(r)}" for r in rows]
|
||||||
return f"Found {len(rows)} note(s):\n" + "\n".join(lines)
|
return f"Found {len(rows)} note(s):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
@@ -72,14 +79,10 @@ async def create_note(
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
row = result["row"]
|
row = result["row"]
|
||||||
# Index the note content in the vector store.
|
note_id: str = row["id"]
|
||||||
vector = await embed(content)
|
# Generate summary asynchronously — fire-and-forget.
|
||||||
await execute_on_client(
|
asyncio.create_task(_refresh_summary(note_id, title, content))
|
||||||
action="vector_upsert",
|
return f"Note created: '{row['title']}' (id: {note_id})."
|
||||||
data={"id": row["id"], "projectId": row.get("projectId"), "content": content},
|
|
||||||
vector=vector,
|
|
||||||
)
|
|
||||||
return f"Note created: '{row['title']}' (id: {row['id']})."
|
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
@@ -88,7 +91,8 @@ async def update_note(
|
|||||||
title: str = "",
|
title: str = "",
|
||||||
content: str = "",
|
content: str = "",
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Update an existing note. Only pass fields that should change.
|
"""Update an existing note directly (no approval required).
|
||||||
|
Use propose_note_edit instead when human review is needed.
|
||||||
note_id: UUID of the note (required)
|
note_id: UUID of the note (required)
|
||||||
If you need to preserve existing content, call get_note first.
|
If you need to preserve existing content, call get_note first.
|
||||||
"""
|
"""
|
||||||
@@ -103,17 +107,63 @@ async def update_note(
|
|||||||
data={"id": note_id, "updates": updates},
|
data={"id": note_id, "updates": updates},
|
||||||
)
|
)
|
||||||
row = result["row"]
|
row = result["row"]
|
||||||
# Re-index if content changed.
|
|
||||||
if content:
|
if content:
|
||||||
vector = await embed(content)
|
new_title = title or row.get("title", "")
|
||||||
await execute_on_client(
|
asyncio.create_task(_refresh_summary(note_id, new_title, content))
|
||||||
action="vector_upsert",
|
|
||||||
data={"id": note_id, "projectId": row.get("projectId"), "content": content},
|
|
||||||
vector=vector,
|
|
||||||
)
|
|
||||||
return f"Note updated: '{row['title']}' (id: {row['id']})."
|
return f"Note updated: '{row['title']}' (id: {row['id']})."
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def propose_note_edit(
|
||||||
|
note_id: str,
|
||||||
|
edit_type: str,
|
||||||
|
proposed_content: str,
|
||||||
|
reasoning: str = "",
|
||||||
|
anchor_before: str = "",
|
||||||
|
anchor_text: str = "",
|
||||||
|
agent_id: str = "",
|
||||||
|
run_id: str = "",
|
||||||
|
) -> str:
|
||||||
|
"""Propose an AI edit to an existing note, pending human approval.
|
||||||
|
|
||||||
|
Use this instead of update_note when review_required is true.
|
||||||
|
The user will see the proposal highlighted before it is merged.
|
||||||
|
|
||||||
|
note_id: UUID of the target note (required)
|
||||||
|
edit_type: 'append' | 'insert' | 'replace'
|
||||||
|
- append: adds proposed_content at the end of the note
|
||||||
|
- insert: inserts proposed_content immediately after anchor_before text
|
||||||
|
- replace: replaces the first occurrence of anchor_text with proposed_content
|
||||||
|
proposed_content: the new Markdown text to add or substitute (required)
|
||||||
|
reasoning: brief explanation shown to the user (recommended)
|
||||||
|
anchor_before: for 'insert' — the text snippet that precedes the insertion point
|
||||||
|
anchor_text: for 'replace' — the exact text to be replaced
|
||||||
|
agent_id: agent identifier (for traceability)
|
||||||
|
run_id: run identifier (for traceability)
|
||||||
|
"""
|
||||||
|
if edit_type not in ("append", "insert", "replace"):
|
||||||
|
return f"Invalid edit_type '{edit_type}'. Use 'append', 'insert', or 'replace'."
|
||||||
|
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="propose_note_edit",
|
||||||
|
data={
|
||||||
|
"noteId": note_id,
|
||||||
|
"type": edit_type,
|
||||||
|
"proposedContent": proposed_content,
|
||||||
|
"reasoning": reasoning or None,
|
||||||
|
"anchorBefore": anchor_before or None,
|
||||||
|
"anchorText": anchor_text or None,
|
||||||
|
"agentId": agent_id or None,
|
||||||
|
"runId": run_id or None,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
edit_id = result.get("id", "?")
|
||||||
|
return (
|
||||||
|
f"Edit proposal created (id: {edit_id}) for note {note_id}. "
|
||||||
|
f"Status: pending user approval."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
async def delete_note(note_id: str) -> str:
|
async def delete_note(note_id: str) -> str:
|
||||||
"""Delete a note permanently by its UUID."""
|
"""Delete a note permanently by its UUID."""
|
||||||
@@ -121,23 +171,36 @@ async def delete_note(note_id: str) -> str:
|
|||||||
return f"Note {note_id} deleted."
|
return f"Note {note_id} deleted."
|
||||||
|
|
||||||
|
|
||||||
@registry.register
|
async def _refresh_summary(note_id: str, title: str, content: str) -> None:
|
||||||
class NoteAgent(ChatAgent):
|
"""Generate and persist the AI summary for a note. Fire-and-forget."""
|
||||||
def get_name(self) -> str:
|
try:
|
||||||
return "note_agent"
|
summary = await generate_note_summary(title, content)
|
||||||
|
if summary:
|
||||||
|
await execute_on_client(
|
||||||
|
action="update",
|
||||||
|
table="notes",
|
||||||
|
data={
|
||||||
|
"id": note_id,
|
||||||
|
"updates": {
|
||||||
|
"aiSummary": summary,
|
||||||
|
"aiSummaryUpdatedAt": int(__import__("time").time() * 1000),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
pass # fire-and-forget; errors logged by generate_note_summary
|
||||||
|
|
||||||
def get_description(self) -> str:
|
|
||||||
return "Manages notes: list, get, create, update, delete"
|
|
||||||
|
|
||||||
def get_tools(self) -> list[Any]:
|
NOTE_TOOLS: list[Any] = [
|
||||||
return [list_notes, get_note, create_note, update_note, delete_note]
|
list_notes,
|
||||||
|
get_note,
|
||||||
|
create_note,
|
||||||
|
update_note,
|
||||||
|
propose_note_edit,
|
||||||
|
delete_note,
|
||||||
|
]
|
||||||
|
|
||||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
NOTE_READ_TOOLS: list[Any] = [
|
||||||
llm = get_llm()
|
list_notes,
|
||||||
messages = [
|
get_note,
|
||||||
SystemMessage(content=_SYSTEM_PROMPT),
|
]
|
||||||
HumanMessage(
|
|
||||||
content=f"User query: {query}\nContext: {json.dumps(context)[:1000]}"
|
|
||||||
),
|
|
||||||
]
|
|
||||||
return await self._tool_loop(llm, messages, self.get_tools())
|
|
||||||
|
|||||||
@@ -4,29 +4,10 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langchain_core.messages import HumanMessage, SystemMessage
|
|
||||||
from langchain_core.tools import tool
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
from app.core.agent_registry import ChatAgent, registry
|
|
||||||
from app.core.llm import get_llm
|
|
||||||
from app.core.ws_context import execute_on_client
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
_SYSTEM_PROMPT = (
|
|
||||||
"You are a project management assistant. You help users create, find,\n"
|
|
||||||
"update, and archive projects in their workspace.\n\n"
|
|
||||||
"Rules:\n"
|
|
||||||
" - status must be one of: active, archived\n"
|
|
||||||
" - client_id is optional; link to a client only when explicitly mentioned\n"
|
|
||||||
" - ai_summary is populated only when the user asks for a project summary;\n"
|
|
||||||
" derive it from context data — do not fabricate content\n"
|
|
||||||
" - Use list_projects for scoped queries; list_all_projects only when the\n"
|
|
||||||
" user wants a complete cross-client view including archived projects\n"
|
|
||||||
" - get_project requires a project UUID; resolve the ID first by calling\n"
|
|
||||||
" list_projects if you only have a project name\n"
|
|
||||||
" - Prefer archiving (update_project status=archived) over deletion;\n"
|
|
||||||
" only call delete_project when the user explicitly confirms deletion."
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
async def list_projects(
|
async def list_projects(
|
||||||
@@ -136,30 +117,17 @@ async def delete_project(project_id: str) -> str:
|
|||||||
return f"Project {project_id} permanently deleted."
|
return f"Project {project_id} permanently deleted."
|
||||||
|
|
||||||
|
|
||||||
@registry.register
|
PROJECT_TOOLS: list[Any] = [
|
||||||
class ProjectAgent(ChatAgent):
|
|
||||||
def get_name(self) -> str:
|
|
||||||
return "project_agent"
|
|
||||||
|
|
||||||
def get_description(self) -> str:
|
|
||||||
return "Manages projects: list, get, create, update, archive, delete"
|
|
||||||
|
|
||||||
def get_tools(self) -> list[Any]:
|
|
||||||
return [
|
|
||||||
list_projects,
|
list_projects,
|
||||||
list_all_projects,
|
list_all_projects,
|
||||||
get_project,
|
get_project,
|
||||||
create_project,
|
create_project,
|
||||||
update_project,
|
update_project,
|
||||||
delete_project,
|
delete_project,
|
||||||
]
|
]
|
||||||
|
|
||||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
PROJECT_READ_TOOLS: list[Any] = [
|
||||||
llm = get_llm()
|
list_projects,
|
||||||
messages = [
|
list_all_projects,
|
||||||
SystemMessage(content=_SYSTEM_PROMPT),
|
get_project,
|
||||||
HumanMessage(
|
]
|
||||||
content=f"User query: {query}\nContext: {json.dumps(context)[:1000]}"
|
|
||||||
),
|
|
||||||
]
|
|
||||||
return await self._tool_loop(llm, messages, self.get_tools())
|
|
||||||
|
|||||||
63
app/agents/relations_agent.py
Normal file
63
app/agents/relations_agent.py
Normal file
@@ -0,0 +1,63 @@
|
|||||||
|
"""Relations agent — read-only tool wrapping MemoryMiddleware.query_relations."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware
|
||||||
|
from app.db import async_session
|
||||||
|
|
||||||
|
# Injected at tool-factory time by _brief_research_tools(); not a module-level global.
|
||||||
|
# Each tool closure captures the user_id bound at factory time.
|
||||||
|
|
||||||
|
|
||||||
|
def make_query_relations_tool(user_id: str, trace_id: str | None = None) -> Any:
|
||||||
|
"""Return a query_relations tool bound to *user_id*."""
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def query_relations(
|
||||||
|
subject_label: str = "",
|
||||||
|
predicate: str = "",
|
||||||
|
object_label: str = "",
|
||||||
|
limit: int = 10,
|
||||||
|
) -> str:
|
||||||
|
"""Query the relational memory graph for entity relationships.
|
||||||
|
|
||||||
|
Returns rows where subject ↔ predicate ↔ object match the given filters.
|
||||||
|
All parameters are optional — omit to retrieve all relations up to limit.
|
||||||
|
|
||||||
|
subject_label: entity label on the left side (e.g. a client name, "Acme Corp").
|
||||||
|
predicate: relationship type (e.g. "mentioned_in", "works_at", "related_to").
|
||||||
|
object_label: entity label on the right side (e.g. a project name, "Website Redesign").
|
||||||
|
limit: max rows to return (default 10).
|
||||||
|
"""
|
||||||
|
import logging
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
logger.info(
|
||||||
|
"relations_agent: query_relations trace=%s user=%s subject=%r predicate=%r object=%r",
|
||||||
|
trace_id or "-", user_id, subject_label, predicate, object_label,
|
||||||
|
)
|
||||||
|
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
rows = await memory.query_relations(
|
||||||
|
user_id=user_id,
|
||||||
|
subject=subject_label or None,
|
||||||
|
predicate=predicate or None,
|
||||||
|
object_=object_label or None,
|
||||||
|
limit=limit,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not rows:
|
||||||
|
return "No relational memory entries found for the given filters."
|
||||||
|
|
||||||
|
lines = [
|
||||||
|
f"- {r.subject_label} —[{r.predicate}]→ {r.object_label}"
|
||||||
|
+ (f" (confidence: {r.confidence:.2f})" if r.confidence is not None else "")
|
||||||
|
for r in rows
|
||||||
|
]
|
||||||
|
return f"Found {len(rows)} relation(s):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
return query_relations
|
||||||
@@ -3,33 +3,22 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
|
import re
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langchain_core.messages import HumanMessage, SystemMessage
|
|
||||||
from langchain_core.tools import tool
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
from app.core.agent_registry import ChatAgent, registry
|
|
||||||
from app.core.llm import get_llm
|
|
||||||
from app.core.ws_context import execute_on_client
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
_SYSTEM_PROMPT = (
|
_UUID_RE = re.compile(
|
||||||
"You are a task management assistant for a project workspace.\n"
|
r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$"
|
||||||
"You create, update, list, and track tasks and their comments.\n\n"
|
|
||||||
"Rules:\n"
|
|
||||||
" - status must be one of: todo, in_progress, done\n"
|
|
||||||
" - priority must be one of: high, medium, low\n"
|
|
||||||
" - due_date is a Unix timestamp in milliseconds; convert human dates\n"
|
|
||||||
" - assignees is a JSON-encoded array of strings (e.g. '[\"Alice\",\"Bob\"]')\n"
|
|
||||||
" - project_id is optional; link to a project when the user mentions one\n"
|
|
||||||
" - is_ai_suggested: 1 only when proactively proposing a task the user\n"
|
|
||||||
" did not explicitly request; 0 otherwise\n"
|
|
||||||
" - is_approved defaults to 0; set to 1 only when the user confirms\n"
|
|
||||||
" - Use list_tasks_due_today for 'what's due today' queries\n"
|
|
||||||
" - For update_task, use -1 for integer fields you do not want to change\n"
|
|
||||||
" - Always confirm the action in plain, user-friendly language."
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_uuid(value: str) -> bool:
|
||||||
|
return bool(_UUID_RE.match(value))
|
||||||
|
|
||||||
|
|
||||||
# ── Task tools ────────────────────────────────────────────────────────
|
# ── Task tools ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@@ -37,31 +26,141 @@ _SYSTEM_PROMPT = (
|
|||||||
async def list_tasks(
|
async def list_tasks(
|
||||||
project_id: str = "",
|
project_id: str = "",
|
||||||
status: str = "",
|
status: str = "",
|
||||||
|
priority: str = "",
|
||||||
|
assignee: str = "",
|
||||||
search: str = "",
|
search: str = "",
|
||||||
order_by: str = "",
|
order_by: str = "",
|
||||||
|
order_dir: str = "",
|
||||||
|
due_date_from: int = -1,
|
||||||
|
due_date_to: int = -1,
|
||||||
|
created_at_from: int = -1,
|
||||||
|
created_at_to: int = -1,
|
||||||
|
completed_at_from: int = -1,
|
||||||
|
completed_at_to: int = -1,
|
||||||
|
is_ai_suggested: int = -1,
|
||||||
|
limit: int = 50,
|
||||||
|
offset: int = 0,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""List tasks, optionally filtered by project_id, status (todo|in_progress|done),
|
"""List tasks with optional filters. Returns up to `limit` results (default 50).
|
||||||
a search string, or an order_by field name (dueDate|priority|createdAt)."""
|
|
||||||
result = await execute_on_client(
|
project_id: UUID of the project to scope results to.
|
||||||
action="select",
|
status: filter by status — todo | in_progress | done.
|
||||||
table="tasks",
|
priority: filter by priority — high | medium | low.
|
||||||
filters={
|
assignee: substring to match against assignee names. OMIT unless the user explicitly
|
||||||
"projectId": project_id or None,
|
names a person or refers to themselves ("my tasks", "assigned to me", "mine").
|
||||||
|
Do NOT default to the current user.
|
||||||
|
search: substring search across title and description.
|
||||||
|
order_by: sort field — dueDate | priority | createdAt | completedAt.
|
||||||
|
order_dir: asc (default) | desc.
|
||||||
|
due_date_from / due_date_to: ms epoch range for dueDate. Use -1 to omit.
|
||||||
|
created_at_from / created_at_to: ms epoch range for createdAt. Use -1 to omit.
|
||||||
|
completed_at_from / completed_at_to: ms epoch range for completedAt. Use -1 to omit.
|
||||||
|
is_ai_suggested: 0 or 1 to filter by AI-suggested flag; -1 = any.
|
||||||
|
limit: max rows to return (default 50). Use with offset to paginate.
|
||||||
|
offset: skip first N rows (default 0).
|
||||||
|
|
||||||
|
Tip — combine *_from and *_to for a closed range; pass only one for open-ended.
|
||||||
|
Tip — prefer count_tasks for "how many" questions to avoid listing rows.
|
||||||
|
Tip — for natural-language windows ("today", "tomorrow", "this week", "last month", etc.)
|
||||||
|
take due_date_from / due_date_to verbatim from the DATE CONTEXT block in the system prompt;
|
||||||
|
do not compute boundaries from the current UTC instant.
|
||||||
|
"""
|
||||||
|
normalized_project_id = project_id if (project_id and _is_uuid(project_id)) else ""
|
||||||
|
filters: dict[str, Any] = {
|
||||||
|
"projectId": normalized_project_id or None,
|
||||||
"status": status or None,
|
"status": status or None,
|
||||||
|
"priority": priority or None,
|
||||||
"search": search or None,
|
"search": search or None,
|
||||||
"orderBy": order_by or None,
|
"orderBy": order_by or None,
|
||||||
},
|
"orderDir": order_dir or None,
|
||||||
)
|
"limit": limit,
|
||||||
|
"offset": offset,
|
||||||
|
}
|
||||||
|
if assignee:
|
||||||
|
filters["assignee"] = assignee
|
||||||
|
if due_date_from != -1:
|
||||||
|
filters["dueDateFrom"] = due_date_from
|
||||||
|
if due_date_to != -1:
|
||||||
|
filters["dueDateTo"] = due_date_to
|
||||||
|
if created_at_from != -1:
|
||||||
|
filters["createdAtFrom"] = created_at_from
|
||||||
|
if created_at_to != -1:
|
||||||
|
filters["createdAtTo"] = created_at_to
|
||||||
|
if completed_at_from != -1:
|
||||||
|
filters["completedAtFrom"] = completed_at_from
|
||||||
|
if completed_at_to != -1:
|
||||||
|
filters["completedAtTo"] = completed_at_to
|
||||||
|
if is_ai_suggested != -1:
|
||||||
|
filters["isAiSuggested"] = is_ai_suggested
|
||||||
|
|
||||||
|
result = await execute_on_client(action="select", table="tasks", filters=filters)
|
||||||
rows = result.get("rows", [])
|
rows = result.get("rows", [])
|
||||||
if not rows:
|
if not rows:
|
||||||
return "No tasks found matching the given filters."
|
return "No tasks found matching the given filters."
|
||||||
lines = [
|
lines = [
|
||||||
f"- {r['title']} (status: {r['status']}, priority: {r['priority']}, id: {r['id']})"
|
f"- {r['title']} (status: {r['status']}, priority: {r['priority']}, "
|
||||||
|
f"dueDate: {r.get('dueDate')}, completedAt: {r.get('completedAt')}, "
|
||||||
|
f"projectId: {r.get('projectId')}, id: {r['id']})"
|
||||||
for r in rows
|
for r in rows
|
||||||
]
|
]
|
||||||
return f"Found {len(rows)} task(s):\n" + "\n".join(lines)
|
return f"Found {len(rows)} task(s):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def count_tasks(
|
||||||
|
project_id: str = "",
|
||||||
|
status: str = "",
|
||||||
|
priority: str = "",
|
||||||
|
assignee: str = "",
|
||||||
|
search: str = "",
|
||||||
|
due_date_from: int = -1,
|
||||||
|
due_date_to: int = -1,
|
||||||
|
created_at_from: int = -1,
|
||||||
|
created_at_to: int = -1,
|
||||||
|
completed_at_from: int = -1,
|
||||||
|
completed_at_to: int = -1,
|
||||||
|
is_ai_suggested: int = -1,
|
||||||
|
) -> str:
|
||||||
|
"""Count tasks matching the given filters without returning rows.
|
||||||
|
|
||||||
|
Use this instead of list_tasks for "how many" questions — it is much cheaper.
|
||||||
|
Same filter parameters as list_tasks (no limit/offset/order_by needed).
|
||||||
|
assignee: OMIT unless the user explicitly names a person or refers to themselves
|
||||||
|
("my tasks"). Do NOT default to the current user.
|
||||||
|
due_date_from / due_date_to: ms epoch range for dueDate. Use -1 to omit.
|
||||||
|
created_at_from / created_at_to: ms epoch range for createdAt. Use -1 to omit.
|
||||||
|
completed_at_from / completed_at_to: ms epoch range for completedAt. Use -1 to omit.
|
||||||
|
Tip — for natural-language windows take due_date_from / due_date_to from the DATE CONTEXT block;
|
||||||
|
do not compute boundaries from the current UTC instant.
|
||||||
|
"""
|
||||||
|
normalized_project_id = project_id if (project_id and _is_uuid(project_id)) else ""
|
||||||
|
filters: dict[str, Any] = {
|
||||||
|
"projectId": normalized_project_id or None,
|
||||||
|
"status": status or None,
|
||||||
|
"priority": priority or None,
|
||||||
|
"search": search or None,
|
||||||
|
}
|
||||||
|
if assignee:
|
||||||
|
filters["assignee"] = assignee
|
||||||
|
if due_date_from != -1:
|
||||||
|
filters["dueDateFrom"] = due_date_from
|
||||||
|
if due_date_to != -1:
|
||||||
|
filters["dueDateTo"] = due_date_to
|
||||||
|
if created_at_from != -1:
|
||||||
|
filters["createdAtFrom"] = created_at_from
|
||||||
|
if created_at_to != -1:
|
||||||
|
filters["createdAtTo"] = created_at_to
|
||||||
|
if completed_at_from != -1:
|
||||||
|
filters["completedAtFrom"] = completed_at_from
|
||||||
|
if completed_at_to != -1:
|
||||||
|
filters["completedAtTo"] = completed_at_to
|
||||||
|
if is_ai_suggested != -1:
|
||||||
|
filters["isAiSuggested"] = is_ai_suggested
|
||||||
|
|
||||||
|
result = await execute_on_client(action="count", table="tasks", filters=filters)
|
||||||
|
return f"Task count: {result.get('count', 0)}"
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
async def create_task(
|
async def create_task(
|
||||||
title: str,
|
title: str,
|
||||||
@@ -72,7 +171,6 @@ async def create_task(
|
|||||||
due_date: int = 0,
|
due_date: int = 0,
|
||||||
project_id: str = "",
|
project_id: str = "",
|
||||||
is_ai_suggested: int = 0,
|
is_ai_suggested: int = 0,
|
||||||
is_approved: int = 0,
|
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Create a new task.
|
"""Create a new task.
|
||||||
title: task title (required)
|
title: task title (required)
|
||||||
@@ -83,7 +181,8 @@ async def create_task(
|
|||||||
due_date: Unix timestamp in milliseconds; 0 means no due date
|
due_date: Unix timestamp in milliseconds; 0 means no due date
|
||||||
project_id: optional UUID of the parent project
|
project_id: optional UUID of the parent project
|
||||||
is_ai_suggested: 1 if proactively suggested, 0 if user-requested
|
is_ai_suggested: 1 if proactively suggested, 0 if user-requested
|
||||||
is_approved: 0 until the user confirms; 1 when confirmed
|
|
||||||
|
completedAt is set automatically when status is 'done'.
|
||||||
"""
|
"""
|
||||||
result = await execute_on_client(
|
result = await execute_on_client(
|
||||||
action="insert",
|
action="insert",
|
||||||
@@ -97,13 +196,12 @@ async def create_task(
|
|||||||
"dueDate": due_date or None,
|
"dueDate": due_date or None,
|
||||||
"projectId": project_id or None,
|
"projectId": project_id or None,
|
||||||
"isAiSuggested": is_ai_suggested,
|
"isAiSuggested": is_ai_suggested,
|
||||||
"isApproved": is_approved,
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
row = result["row"]
|
row = result["row"]
|
||||||
return (
|
return (
|
||||||
f"Task created: '{row['title']}' "
|
f"Task created: '{row['title']}' "
|
||||||
f"(id: {row['id']}, status: {row['status']}, priority: {row['priority']})"
|
f"(id: {row['id']}, status: {row['status']}, priority: {row['priority']}, projectId: {row.get('projectId')})"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -117,12 +215,14 @@ async def update_task(
|
|||||||
assignees: str = "",
|
assignees: str = "",
|
||||||
due_date: int = -1,
|
due_date: int = -1,
|
||||||
project_id: str = "",
|
project_id: str = "",
|
||||||
is_approved: int = -1,
|
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Update fields on an existing task. Only pass fields you want to change.
|
"""Update fields on an existing task. Only pass fields you want to change.
|
||||||
task_id: the task's UUID (required)
|
task_id: the task's UUID (required)
|
||||||
due_date: -1 means unchanged; 0 clears the due date; any positive value sets it
|
due_date: -1 means unchanged; 0 clears the due date; any positive value sets it
|
||||||
is_approved: -1 means unchanged; 0 or 1 sets the value
|
|
||||||
|
completedAt is managed automatically:
|
||||||
|
- setting status to 'done' records the current timestamp
|
||||||
|
- changing status away from 'done' clears completedAt
|
||||||
"""
|
"""
|
||||||
updates: dict[str, Any] = {}
|
updates: dict[str, Any] = {}
|
||||||
if title:
|
if title:
|
||||||
@@ -139,15 +239,13 @@ async def update_task(
|
|||||||
updates["dueDate"] = due_date or None
|
updates["dueDate"] = due_date or None
|
||||||
if project_id:
|
if project_id:
|
||||||
updates["projectId"] = project_id
|
updates["projectId"] = project_id
|
||||||
if is_approved != -1:
|
|
||||||
updates["isApproved"] = is_approved
|
|
||||||
result = await execute_on_client(
|
result = await execute_on_client(
|
||||||
action="update",
|
action="update",
|
||||||
table="tasks",
|
table="tasks",
|
||||||
data={"id": task_id, "updates": updates},
|
data={"id": task_id, "updates": updates},
|
||||||
)
|
)
|
||||||
row = result["row"]
|
row = result["row"]
|
||||||
return f"Task updated: '{row['title']}' (id: {row['id']}, status: {row['status']})"
|
return f"Task updated: '{row['title']}' (id: {row['id']}, status: {row['status']}, projectId: {row.get('projectId')})"
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
@@ -158,21 +256,36 @@ async def delete_task(task_id: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
async def list_tasks_due_today() -> str:
|
async def list_tasks_due_today(user_timezone: str = "UTC", include_done: bool = False) -> str:
|
||||||
"""List all tasks whose due date falls on today's date."""
|
"""List all tasks whose due date falls on today's date.
|
||||||
now = datetime.now(tz=timezone.utc)
|
|
||||||
start_ms = int(datetime(now.year, now.month, now.day, tzinfo=timezone.utc).timestamp() * 1000)
|
user_timezone: IANA timezone name (e.g. 'Europe/Rome', 'America/New_York').
|
||||||
end_ms = start_ms + 86_400_000 - 1 # last ms of today
|
Always pass the user's timezone so 'today' is computed in their local time.
|
||||||
|
include_done: set True to also include already-completed tasks due today (default False).
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from zoneinfo import ZoneInfo
|
||||||
|
tz = ZoneInfo(user_timezone or "UTC")
|
||||||
|
except Exception:
|
||||||
|
tz = timezone.utc
|
||||||
|
now_local = datetime.now(tz=tz)
|
||||||
|
start_dt = datetime(now_local.year, now_local.month, now_local.day, tzinfo=tz)
|
||||||
|
start_ms = int(start_dt.timestamp() * 1000)
|
||||||
|
end_ms = start_ms + 86_400_000 - 1
|
||||||
|
filters: dict[str, Any] = {"dueDateFrom": start_ms, "dueDateTo": end_ms}
|
||||||
|
if not include_done:
|
||||||
|
filters["status"] = "todo"
|
||||||
result = await execute_on_client(
|
result = await execute_on_client(
|
||||||
action="select",
|
action="select",
|
||||||
table="tasks",
|
table="tasks",
|
||||||
filters={"dueDateFrom": start_ms, "dueDateTo": end_ms},
|
filters=filters,
|
||||||
)
|
)
|
||||||
rows = result.get("rows", [])
|
rows = result.get("rows", [])
|
||||||
if not rows:
|
if not rows:
|
||||||
return "No tasks are due today."
|
return "No tasks are due today."
|
||||||
lines = [
|
lines = [
|
||||||
f"- {r['title']} (priority: {r['priority']}, status: {r['status']}, id: {r['id']})"
|
f"- {r['title']} (priority: {r['priority']}, status: {r['status']}, "
|
||||||
|
f"projectId: {r.get('projectId')}, id: {r['id']})"
|
||||||
for r in rows
|
for r in rows
|
||||||
]
|
]
|
||||||
return f"Tasks due today ({len(rows)}):\n" + "\n".join(lines)
|
return f"Tasks due today ({len(rows)}):\n" + "\n".join(lines)
|
||||||
@@ -208,8 +321,11 @@ async def add_task_comment(task_id: str, author: str, content: str) -> str:
|
|||||||
table="taskComments",
|
table="taskComments",
|
||||||
data={"taskId": task_id, "author": author, "content": content},
|
data={"taskId": task_id, "author": author, "content": content},
|
||||||
)
|
)
|
||||||
row = result["row"]
|
row = result.get("row", {})
|
||||||
return f"Comment added by {row['author']} on task {row['taskId']} (comment id: {row['id']})."
|
row_author = row.get("author", author)
|
||||||
|
row_task_id = row.get("taskId") or row.get("task_id") or task_id
|
||||||
|
row_comment_id = row.get("id", "unknown")
|
||||||
|
return f"Comment added by {row_author} on task {row_task_id} (comment id: {row_comment_id})."
|
||||||
|
|
||||||
|
|
||||||
@tool
|
@tool
|
||||||
@@ -222,17 +338,9 @@ async def delete_task_comment(comment_id: str) -> str:
|
|||||||
# ── Agent ─────────────────────────────────────────────────────────────
|
# ── Agent ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@registry.register
|
TASK_TOOLS: list[Any] = [
|
||||||
class TaskAgent(ChatAgent):
|
|
||||||
def get_name(self) -> str:
|
|
||||||
return "task_agent"
|
|
||||||
|
|
||||||
def get_description(self) -> str:
|
|
||||||
return "Manages tasks and comments: list, create, update, delete, due-today, comments"
|
|
||||||
|
|
||||||
def get_tools(self) -> list[Any]:
|
|
||||||
return [
|
|
||||||
list_tasks,
|
list_tasks,
|
||||||
|
count_tasks,
|
||||||
create_task,
|
create_task,
|
||||||
update_task,
|
update_task,
|
||||||
delete_task,
|
delete_task,
|
||||||
@@ -240,14 +348,11 @@ class TaskAgent(ChatAgent):
|
|||||||
list_task_comments,
|
list_task_comments,
|
||||||
add_task_comment,
|
add_task_comment,
|
||||||
delete_task_comment,
|
delete_task_comment,
|
||||||
]
|
]
|
||||||
|
|
||||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
TASK_READ_TOOLS: list[Any] = [
|
||||||
llm = get_llm()
|
list_tasks,
|
||||||
messages = [
|
count_tasks,
|
||||||
SystemMessage(content=_SYSTEM_PROMPT),
|
list_tasks_due_today,
|
||||||
HumanMessage(
|
list_task_comments,
|
||||||
content=f"User query: {query}\nContext: {json.dumps(context)[:1000]}"
|
]
|
||||||
),
|
|
||||||
]
|
|
||||||
return await self._tool_loop(llm, messages, self.get_tools())
|
|
||||||
|
|||||||
270
app/agents/timeline_agent.py
Normal file
270
app/agents/timeline_agent.py
Normal file
@@ -0,0 +1,270 @@
|
|||||||
|
"""Timeline agent — project milestone management (list, create, update, delete)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
|
from app.core.ws_context import execute_on_client
|
||||||
|
|
||||||
|
_UUID_RE = re.compile(
|
||||||
|
r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_uuid(value: str) -> bool:
|
||||||
|
return bool(_UUID_RE.match(value))
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def list_timelines(
|
||||||
|
project_id: str = "",
|
||||||
|
type: str = "",
|
||||||
|
is_completed: int = -1,
|
||||||
|
is_ai_suggested: int = -1,
|
||||||
|
order_by: str = "",
|
||||||
|
order_dir: str = "",
|
||||||
|
date_from: int = -1,
|
||||||
|
date_to: int = -1,
|
||||||
|
created_at_from: int = -1,
|
||||||
|
created_at_to: int = -1,
|
||||||
|
completed_at_from: int = -1,
|
||||||
|
completed_at_to: int = -1,
|
||||||
|
limit: int = 50,
|
||||||
|
offset: int = 0,
|
||||||
|
) -> str:
|
||||||
|
"""List timeline events (milestones, checkpoints, activities) with optional filters.
|
||||||
|
|
||||||
|
project_id: UUID to scope results to a specific project.
|
||||||
|
type: filter by event type — milestone | checkpoint | activity.
|
||||||
|
is_completed: 0 = incomplete only, 1 = completed only, -1 = any (default).
|
||||||
|
is_ai_suggested: 0 or 1 to filter by AI-suggested flag; -1 = any.
|
||||||
|
order_by: sort field — date (default) | createdAt | completedAt.
|
||||||
|
order_dir: asc (default) | desc.
|
||||||
|
date_from / date_to: ms epoch range for the event date. Use -1 to omit.
|
||||||
|
created_at_from / created_at_to: ms epoch range for createdAt. Use -1 to omit.
|
||||||
|
completed_at_from / completed_at_to: ms epoch range for completedAt. Use -1 to omit.
|
||||||
|
limit: max rows to return (default 50). Use with offset to paginate.
|
||||||
|
offset: skip first N rows (default 0).
|
||||||
|
|
||||||
|
Tip — combine *_from and *_to for a closed range; pass only one for open-ended.
|
||||||
|
Tip — prefer count_timelines for "how many" questions to avoid listing rows.
|
||||||
|
Tip — for natural-language windows ("today", "this week", "last month", etc.)
|
||||||
|
take date_from / date_to verbatim from the DATE CONTEXT block in the system prompt;
|
||||||
|
do not compute boundaries from the current UTC instant.
|
||||||
|
"""
|
||||||
|
normalized_project_id = project_id if (project_id and _is_uuid(project_id)) else ""
|
||||||
|
filters: dict[str, Any] = {
|
||||||
|
"projectId": normalized_project_id or None,
|
||||||
|
"orderBy": order_by or None,
|
||||||
|
"orderDir": order_dir or None,
|
||||||
|
"limit": limit,
|
||||||
|
"offset": offset,
|
||||||
|
}
|
||||||
|
if type:
|
||||||
|
filters["type"] = type
|
||||||
|
if is_completed != -1:
|
||||||
|
filters["isCompleted"] = is_completed
|
||||||
|
if is_ai_suggested != -1:
|
||||||
|
filters["isAiSuggested"] = is_ai_suggested
|
||||||
|
if date_from != -1:
|
||||||
|
filters["dateFrom"] = date_from
|
||||||
|
if date_to != -1:
|
||||||
|
filters["dateTo"] = date_to
|
||||||
|
if created_at_from != -1:
|
||||||
|
filters["createdAtFrom"] = created_at_from
|
||||||
|
if created_at_to != -1:
|
||||||
|
filters["createdAtTo"] = created_at_to
|
||||||
|
if completed_at_from != -1:
|
||||||
|
filters["completedAtFrom"] = completed_at_from
|
||||||
|
if completed_at_to != -1:
|
||||||
|
filters["completedAtTo"] = completed_at_to
|
||||||
|
|
||||||
|
result = await execute_on_client(action="select", table="timelines", filters=filters)
|
||||||
|
rows = result.get("rows", [])
|
||||||
|
if not rows:
|
||||||
|
return "No timeline events found."
|
||||||
|
lines = [
|
||||||
|
f"- {r['title']} (date: {r['date']}, type: {r.get('type')}, "
|
||||||
|
f"completed: {bool(r.get('isCompleted'))}, completedAt: {r.get('completedAt')}, "
|
||||||
|
f"projectId: {r.get('projectId')}, id: {r['id']})"
|
||||||
|
for r in rows
|
||||||
|
]
|
||||||
|
return f"Found {len(rows)} timeline event(s):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def count_timelines(
|
||||||
|
project_id: str = "",
|
||||||
|
type: str = "",
|
||||||
|
is_completed: int = -1,
|
||||||
|
is_ai_suggested: int = -1,
|
||||||
|
date_from: int = -1,
|
||||||
|
date_to: int = -1,
|
||||||
|
created_at_from: int = -1,
|
||||||
|
created_at_to: int = -1,
|
||||||
|
completed_at_from: int = -1,
|
||||||
|
completed_at_to: int = -1,
|
||||||
|
) -> str:
|
||||||
|
"""Count timeline events matching the given filters without returning rows.
|
||||||
|
|
||||||
|
Use this instead of list_timelines for "how many" questions — it is much cheaper.
|
||||||
|
Same filter parameters as list_timelines (no limit/offset/order_by needed).
|
||||||
|
|
||||||
|
date_from / date_to: ms epoch range for the event date. Use -1 to omit.
|
||||||
|
completed_at_from / completed_at_to: ms epoch range for completedAt. Use -1 to omit.
|
||||||
|
Tip — for natural-language windows take date_from / date_to from the DATE CONTEXT block;
|
||||||
|
do not compute boundaries from the current UTC instant.
|
||||||
|
"""
|
||||||
|
normalized_project_id = project_id if (project_id and _is_uuid(project_id)) else ""
|
||||||
|
filters: dict[str, Any] = {"projectId": normalized_project_id or None}
|
||||||
|
if type:
|
||||||
|
filters["type"] = type
|
||||||
|
if is_completed != -1:
|
||||||
|
filters["isCompleted"] = is_completed
|
||||||
|
if is_ai_suggested != -1:
|
||||||
|
filters["isAiSuggested"] = is_ai_suggested
|
||||||
|
if date_from != -1:
|
||||||
|
filters["dateFrom"] = date_from
|
||||||
|
if date_to != -1:
|
||||||
|
filters["dateTo"] = date_to
|
||||||
|
if created_at_from != -1:
|
||||||
|
filters["createdAtFrom"] = created_at_from
|
||||||
|
if created_at_to != -1:
|
||||||
|
filters["createdAtTo"] = created_at_to
|
||||||
|
if completed_at_from != -1:
|
||||||
|
filters["completedAtFrom"] = completed_at_from
|
||||||
|
if completed_at_to != -1:
|
||||||
|
filters["completedAtTo"] = completed_at_to
|
||||||
|
|
||||||
|
result = await execute_on_client(action="count", table="timelines", filters=filters)
|
||||||
|
return f"Timeline event count: {result.get('count', 0)}"
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def create_timeline(
|
||||||
|
project_id: str,
|
||||||
|
title: str,
|
||||||
|
date: int,
|
||||||
|
type: str = "milestone",
|
||||||
|
is_completed: int = 0,
|
||||||
|
is_ai_suggested: int = 0,
|
||||||
|
) -> str:
|
||||||
|
"""Create a project timeline event.
|
||||||
|
project_id: REQUIRED UUID of the parent project
|
||||||
|
title: descriptive name for the event
|
||||||
|
date: Unix timestamp in milliseconds for the event date
|
||||||
|
type: milestone (default) | checkpoint | activity
|
||||||
|
is_completed: 1 if already completed, 0 if not (default 0)
|
||||||
|
is_ai_suggested: 1 if proactively suggested, 0 if user-requested
|
||||||
|
|
||||||
|
completedAt is set automatically when is_completed is 1.
|
||||||
|
"""
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="insert",
|
||||||
|
table="timelines",
|
||||||
|
data={
|
||||||
|
"projectId": project_id,
|
||||||
|
"title": title,
|
||||||
|
"date": date,
|
||||||
|
"type": type,
|
||||||
|
"isCompleted": is_completed,
|
||||||
|
"isAiSuggested": is_ai_suggested,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
row = result["row"]
|
||||||
|
return f"Timeline event created: '{row['title']}' (id: {row['id']}, date: {row['date']}, type: {row.get('type')})"
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def update_timeline(
|
||||||
|
timeline_id: str,
|
||||||
|
title: str = "",
|
||||||
|
date: int = -1,
|
||||||
|
is_completed: int = -1,
|
||||||
|
) -> str:
|
||||||
|
"""Update a timeline event. Only pass fields that should change.
|
||||||
|
timeline_id: UUID of the event (required)
|
||||||
|
date: -1 means unchanged; any other value sets the new date (ms timestamp)
|
||||||
|
is_completed: 0 = mark incomplete, 1 = mark complete, -1 = unchanged
|
||||||
|
|
||||||
|
completedAt is managed automatically:
|
||||||
|
- setting is_completed to 1 records the current timestamp
|
||||||
|
- setting is_completed to 0 clears completedAt
|
||||||
|
"""
|
||||||
|
updates: dict[str, Any] = {}
|
||||||
|
if title:
|
||||||
|
updates["title"] = title
|
||||||
|
if date != -1:
|
||||||
|
updates["date"] = date
|
||||||
|
if is_completed != -1:
|
||||||
|
updates["isCompleted"] = is_completed
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="update",
|
||||||
|
table="timelines",
|
||||||
|
data={"id": timeline_id, "updates": updates},
|
||||||
|
)
|
||||||
|
row = result["row"]
|
||||||
|
return f"Timeline event updated: '{row['title']}' (id: {row['id']})"
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def delete_timeline(timeline_id: str) -> str:
|
||||||
|
"""Delete a timeline event permanently by its UUID."""
|
||||||
|
await execute_on_client(action="delete", table="timelines", data={"id": timeline_id})
|
||||||
|
return f"Timeline event {timeline_id} deleted."
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
async def list_timelines_today(user_timezone: str = "UTC", include_completed: bool = True) -> str:
|
||||||
|
"""List all timeline events whose date falls on today.
|
||||||
|
|
||||||
|
user_timezone: IANA timezone name (e.g. 'Europe/Rome', 'America/New_York').
|
||||||
|
Always pass the user's timezone so 'today' is computed in their local time.
|
||||||
|
include_completed: set False to exclude already-completed events (default True).
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from zoneinfo import ZoneInfo
|
||||||
|
tz = ZoneInfo(user_timezone or "UTC")
|
||||||
|
except Exception:
|
||||||
|
tz = timezone.utc
|
||||||
|
now_local = datetime.now(tz=tz)
|
||||||
|
start_dt = datetime(now_local.year, now_local.month, now_local.day, tzinfo=tz)
|
||||||
|
start_ms = int(start_dt.timestamp() * 1000)
|
||||||
|
end_ms = start_ms + 86_400_000 - 1
|
||||||
|
filters: dict[str, Any] = {"dateFrom": start_ms, "dateTo": end_ms}
|
||||||
|
if not include_completed:
|
||||||
|
filters["isCompleted"] = 0
|
||||||
|
result = await execute_on_client(
|
||||||
|
action="select",
|
||||||
|
table="timelines",
|
||||||
|
filters=filters,
|
||||||
|
)
|
||||||
|
rows = result.get("rows", [])
|
||||||
|
if not rows:
|
||||||
|
return "No timeline events today."
|
||||||
|
lines = [
|
||||||
|
f"- {r['title']} (date: {r['date']}, type: {r.get('type')}, "
|
||||||
|
f"completed: {bool(r.get('isCompleted'))}, projectId: {r.get('projectId')}, id: {r['id']})"
|
||||||
|
for r in rows
|
||||||
|
]
|
||||||
|
return f"Timeline events today ({len(rows)}):\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
TIMELINE_TOOLS: list[Any] = [
|
||||||
|
list_timelines,
|
||||||
|
count_timelines,
|
||||||
|
list_timelines_today,
|
||||||
|
create_timeline,
|
||||||
|
update_timeline,
|
||||||
|
delete_timeline,
|
||||||
|
]
|
||||||
|
|
||||||
|
TIMELINE_READ_TOOLS: list[Any] = [
|
||||||
|
list_timelines,
|
||||||
|
count_timelines,
|
||||||
|
list_timelines_today,
|
||||||
|
]
|
||||||
@@ -55,11 +55,49 @@ async def get_current_user(
|
|||||||
raise credentials_exc
|
raise credentials_exc
|
||||||
|
|
||||||
# Live tier lookup — subscription row is the authoritative source.
|
# Live tier lookup — subscription row is the authoritative source.
|
||||||
from app.models import Subscription # noqa: PLC0415
|
# In dev, fall back to 'power' (unlimited) so quota limits don't
|
||||||
|
# block local development when no Stripe subscription exists.
|
||||||
|
from app.models import Subscription, User # noqa: PLC0415
|
||||||
|
|
||||||
result = await db.execute(
|
result = await db.execute(
|
||||||
select(Subscription.tier).where(Subscription.user_id == user_id)
|
select(Subscription.tier).where(Subscription.user_id == user_id)
|
||||||
)
|
)
|
||||||
tier: str = result.scalar_one_or_none() or "free"
|
default_tier = "power" if settings.ENV == "dev" else "free"
|
||||||
|
tier: str = result.scalar_one_or_none() or default_tier
|
||||||
|
|
||||||
return UserProfile(id=user_id, email=email, tier=tier) # type: ignore[arg-type]
|
# Fetch name/surname/avatar_url/onboarding_completed_at/password_hash from user row.
|
||||||
|
user_result = await db.execute(
|
||||||
|
select(
|
||||||
|
User.name, User.surname, User.avatar_url, User.onboarding_completed_at,
|
||||||
|
User.password_hash,
|
||||||
|
).where(User.id == user_id)
|
||||||
|
)
|
||||||
|
user_row = user_result.one_or_none()
|
||||||
|
|
||||||
|
# Convert onboarding_completed_at to epoch ms (int) or None.
|
||||||
|
onboarding_ms: int | None = None
|
||||||
|
if user_row and user_row.onboarding_completed_at is not None:
|
||||||
|
onboarding_ms = int(user_row.onboarding_completed_at.timestamp() * 1000)
|
||||||
|
|
||||||
|
# Load decrypted core memory.
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware # noqa: PLC0415
|
||||||
|
|
||||||
|
memory_dict: dict[str, str] = {}
|
||||||
|
try:
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
blocks = await mw.list_core_blocks(user_id)
|
||||||
|
memory_dict = {b["label"]: b["value"] for b in blocks}
|
||||||
|
except Exception:
|
||||||
|
pass # Non-critical — return empty memory on failure
|
||||||
|
|
||||||
|
return UserProfile(
|
||||||
|
id=user_id,
|
||||||
|
email=email,
|
||||||
|
name=user_row.name if user_row else None,
|
||||||
|
surname=user_row.surname if user_row else None,
|
||||||
|
avatar_url=user_row.avatar_url if user_row else None,
|
||||||
|
has_password=bool(user_row.password_hash) if user_row else False,
|
||||||
|
tier=tier,
|
||||||
|
onboarding_completed_at=onboarding_ms,
|
||||||
|
memory=memory_dict,
|
||||||
|
) # type: ignore[arg-type]
|
||||||
|
|||||||
@@ -8,8 +8,7 @@ that could reveal server-side prompt IP:
|
|||||||
- Internal reasoning markers (<thinking>, <reasoning>, [INST], …)
|
- Internal reasoning markers (<thinking>, <reasoning>, [INST], …)
|
||||||
- Exact-match known prompt fingerprints
|
- Exact-match known prompt fingerprints
|
||||||
|
|
||||||
Binary responses (storage blobs, backup data) are never touched — the
|
The middleware only activates for paths under /api/v1/chat.
|
||||||
middleware only activates for paths under /api/v1/chat.
|
|
||||||
|
|
||||||
Any sanitisation event is logged as a WARNING with the request path and the
|
Any sanitisation event is logged as a WARNING with the request path and the
|
||||||
names of the fields that were modified.
|
names of the fields that were modified.
|
||||||
|
|||||||
513
app/api/routes/agent_setup.py
Normal file
513
app/api/routes/agent_setup.py
Normal file
@@ -0,0 +1,513 @@
|
|||||||
|
"""Chatbot Journey — WS-based guided conversation to build an AgentConfig.
|
||||||
|
|
||||||
|
The journey is driven entirely through WebSocket frames (no REST endpoints).
|
||||||
|
The device WS handler dispatches ``journey_start`` and ``journey_message``
|
||||||
|
frames to the functions exported here.
|
||||||
|
|
||||||
|
Journey flow:
|
||||||
|
1. FE sends ``journey_start`` frame with basic agent info (directory,
|
||||||
|
data_types, schedule).
|
||||||
|
2. Server creates an in-memory session, sets up a WS executor so the
|
||||||
|
setup LLM can use file-system tools, does a first directory scrape,
|
||||||
|
and sends back a ``journey_reply`` with the first question.
|
||||||
|
3. FE sends ``journey_message`` frames for each user reply.
|
||||||
|
4. Server appends the user message, calls the LLM (which may read files
|
||||||
|
via tools), and sends back a ``journey_reply``.
|
||||||
|
5. After 3-5 turns the LLM wraps up by emitting an ``AgentConfig`` JSON
|
||||||
|
block delimited by ``AGENT_CONFIG_START`` / ``AGENT_CONFIG_END``.
|
||||||
|
6. Server parses and validates the JSON with Pydantic, sends
|
||||||
|
``journey_reply`` with ``done=True`` and the serialised config.
|
||||||
|
FE stores it locally.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||||
|
|
||||||
|
from app.agents.filesystem_agent import make_directory_tools
|
||||||
|
from app.core.langfuse_client import compile_prompt, extract_usage, get_langfuse, get_prompt_or_fallback, langfuse_context
|
||||||
|
from app.core.llm import get_agent_llm, model_for_agent
|
||||||
|
from app.schemas import AgentConfig
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ── Session TTL ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
_SESSION_TTL_SECONDS: int = 1800 # 30 minutes
|
||||||
|
|
||||||
|
# Sentinel strings used to delimit the LLM-produced AgentConfig JSON.
|
||||||
|
_CONFIG_START = "AGENT_CONFIG_START"
|
||||||
|
_CONFIG_END = "AGENT_CONFIG_END"
|
||||||
|
|
||||||
|
# Minimum turns before we consider nudging the LLM to wrap up.
|
||||||
|
_MIN_TURNS_BEFORE_NUDGE: int = 3
|
||||||
|
# Hard cap to avoid infinite loops (safety net, not the primary stopping criterion).
|
||||||
|
_MAX_TURNS: int = 15
|
||||||
|
# Max tool-calling steps per LLM invocation.
|
||||||
|
_MAX_TOOL_STEPS: int = 6
|
||||||
|
|
||||||
|
# ── In-memory session store ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class JourneySession:
|
||||||
|
session_id: str
|
||||||
|
user_id: str
|
||||||
|
agent_type: str # "local" | "cloud"
|
||||||
|
directory: str
|
||||||
|
data_types: list[str]
|
||||||
|
history: list[dict[str, Any]] = field(default_factory=list)
|
||||||
|
system_prompt: str = ""
|
||||||
|
langfuse_prompt: Any = None
|
||||||
|
created_at: float = field(default_factory=time.monotonic)
|
||||||
|
|
||||||
|
def is_expired(self) -> bool:
|
||||||
|
return (time.monotonic() - self.created_at) > _SESSION_TTL_SECONDS
|
||||||
|
|
||||||
|
|
||||||
|
# session_id → session
|
||||||
|
_sessions: dict[str, JourneySession] = {}
|
||||||
|
|
||||||
|
|
||||||
|
def get_journey_session(session_id: str, user_id: str) -> JourneySession | None:
|
||||||
|
"""Retrieve session; return None on missing, expired, or wrong owner."""
|
||||||
|
s = _sessions.get(session_id)
|
||||||
|
if s is None or s.is_expired():
|
||||||
|
_sessions.pop(session_id, None)
|
||||||
|
return None
|
||||||
|
if s.user_id != user_id:
|
||||||
|
return None
|
||||||
|
return s
|
||||||
|
|
||||||
|
|
||||||
|
# ── System prompt ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
_JOURNEY_SYSTEM_PROMPT = """\
|
||||||
|
You are a friendly assistant helping a freelancer configure a data-extraction agent.
|
||||||
|
Your job is to understand what files the user has in their directory and produce a
|
||||||
|
structured AgentConfig JSON that the extraction agent will use as its instruction set.
|
||||||
|
|
||||||
|
You have access to file-system tools to explore the user's directory:
|
||||||
|
- list_directory: see folder structure and file names
|
||||||
|
- read_file_content: peek at a file's content
|
||||||
|
- get_file_metadata: check file size, extension, dates
|
||||||
|
|
||||||
|
The user's configured directory is: {directory}
|
||||||
|
Target data types: {data_types}
|
||||||
|
|
||||||
|
## Your process
|
||||||
|
|
||||||
|
### Step 1 — Explore the directory
|
||||||
|
Use list_directory and read_file_content to understand what types of files are present
|
||||||
|
(HTML emails, plain-text documents, CSVs, etc.).
|
||||||
|
|
||||||
|
### Step 2 — Identify content types
|
||||||
|
For each distinct file type found, decide:
|
||||||
|
- A short id (e.g. "email_html", "plain_text", "csv")
|
||||||
|
- Which preprocessing handler to use: "email_html" for HTML emails, "generic" for everything else
|
||||||
|
- A human-readable label and optional detection_hint
|
||||||
|
|
||||||
|
### Step 3 — Ask focused questions (one at a time)
|
||||||
|
Cover these topics based on what you discovered:
|
||||||
|
1. How to map content to entity types (task / note / timeline entry)
|
||||||
|
2. Field mapping rules (e.g. email Subject → task title, filename → note title)
|
||||||
|
3. Priority or status rules (e.g. "urgent" in subject → high priority)
|
||||||
|
4. Date extraction (e.g. "by Friday" → dueDate)
|
||||||
|
5. Exclusion rules (e.g. skip newsletters, skip files with no project match)
|
||||||
|
|
||||||
|
### Step 4 — Produce the AgentConfig JSON
|
||||||
|
Once you are ≥ 90% confident, output the final config between these exact markers
|
||||||
|
(each on its own line):
|
||||||
|
|
||||||
|
{config_start}
|
||||||
|
{{
|
||||||
|
"content_types": [
|
||||||
|
{{
|
||||||
|
"id": "email_html",
|
||||||
|
"label": "Email HTML",
|
||||||
|
"detection_hint": "HTML file with From/To/Subject headers",
|
||||||
|
"preprocessing": "email_html",
|
||||||
|
"extraction_prompt": "Detailed extraction instructions for this content type..."
|
||||||
|
}}
|
||||||
|
],
|
||||||
|
"global_rules": [
|
||||||
|
"If the file cannot be matched to any project, do not create any entity."
|
||||||
|
],
|
||||||
|
"data_types": {data_types_json}
|
||||||
|
}}
|
||||||
|
{config_end}
|
||||||
|
|
||||||
|
## Rules for the extraction_prompt field
|
||||||
|
- Describe when to create a task vs note vs timeline entry (be specific and concrete)
|
||||||
|
- Include field mapping rules based on what you found in the directory
|
||||||
|
- Include priority/status/date rules if applicable
|
||||||
|
- Do NOT include projectId logic — the runner handles project assignment automatically
|
||||||
|
- Do NOT mention isAiSuggested — the runner always sets it to 1
|
||||||
|
|
||||||
|
## Constraints
|
||||||
|
- Never ask about projects, projectId, or how to link records to projects
|
||||||
|
- Never include projectId or project creation logic in the generated config
|
||||||
|
- Keep asking questions until ≥ 90% confident, then output the JSON immediately
|
||||||
|
|
||||||
|
{existing_section}\
|
||||||
|
Begin by exploring the directory, then ask your first question.\
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def _build_system_prompt(
|
||||||
|
directory: str,
|
||||||
|
data_types: list[str],
|
||||||
|
existing_config: str | None = None,
|
||||||
|
) -> tuple[str, Any]:
|
||||||
|
"""Return ``(compiled_system_prompt, langfuse_prompt_obj_or_None)``."""
|
||||||
|
existing_section = (
|
||||||
|
"\nThe user already has the following AgentConfig — refine it based on their answers:\n"
|
||||||
|
f"```json\n{existing_config}\n```\n"
|
||||||
|
if existing_config
|
||||||
|
else ""
|
||||||
|
)
|
||||||
|
template, prompt_obj = get_prompt_or_fallback(
|
||||||
|
"journey_system", _JOURNEY_SYSTEM_PROMPT
|
||||||
|
)
|
||||||
|
compiled = compile_prompt(
|
||||||
|
template,
|
||||||
|
prompt_obj,
|
||||||
|
directory=directory,
|
||||||
|
data_types=", ".join(data_types),
|
||||||
|
data_types_json=json.dumps(data_types),
|
||||||
|
config_start=_CONFIG_START,
|
||||||
|
config_end=_CONFIG_END,
|
||||||
|
existing_section=existing_section,
|
||||||
|
)
|
||||||
|
return compiled, prompt_obj
|
||||||
|
|
||||||
|
|
||||||
|
# ── AgentConfig extraction ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_agent_config(text: str) -> str | None:
|
||||||
|
"""Return validated AgentConfig JSON string from between markers, or None.
|
||||||
|
|
||||||
|
Parses the JSON with Pydantic to ensure it conforms to the schema before
|
||||||
|
returning. Returns None if markers are absent or JSON is invalid.
|
||||||
|
"""
|
||||||
|
if _CONFIG_START not in text or _CONFIG_END not in text:
|
||||||
|
return None
|
||||||
|
start_idx = text.index(_CONFIG_START) + len(_CONFIG_START)
|
||||||
|
end_idx = text.index(_CONFIG_END)
|
||||||
|
raw = text[start_idx:end_idx].strip()
|
||||||
|
if not raw:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
parsed = AgentConfig.model_validate_json(raw)
|
||||||
|
return parsed.model_dump_json()
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("agent_setup: failed to parse AgentConfig JSON: %s", exc)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# ── LLM call with tool support ───────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def _as_text(content: Any) -> str:
|
||||||
|
if content is None:
|
||||||
|
return ""
|
||||||
|
if isinstance(content, str):
|
||||||
|
return content
|
||||||
|
if isinstance(content, list):
|
||||||
|
parts: list[str] = []
|
||||||
|
for item in content:
|
||||||
|
if isinstance(item, str):
|
||||||
|
parts.append(item)
|
||||||
|
elif isinstance(item, dict):
|
||||||
|
text = item.get("text")
|
||||||
|
if isinstance(text, str):
|
||||||
|
parts.append(text)
|
||||||
|
return "".join(parts)
|
||||||
|
return str(content)
|
||||||
|
|
||||||
|
|
||||||
|
async def _call_llm_with_tools(
|
||||||
|
system_prompt: str,
|
||||||
|
history: list[dict[str, Any]],
|
||||||
|
tools: list[Any],
|
||||||
|
*,
|
||||||
|
user_id: str = "",
|
||||||
|
session_id: str = "",
|
||||||
|
langfuse_prompt: Any = None,
|
||||||
|
) -> str:
|
||||||
|
"""Build LangChain messages from history and invoke the LLM with tools.
|
||||||
|
|
||||||
|
Handles tool-calling loops: if the LLM calls tools, execute them and
|
||||||
|
continue until a final text response is produced.
|
||||||
|
"""
|
||||||
|
lf = get_langfuse()
|
||||||
|
messages: list[Any] = [SystemMessage(content=system_prompt)]
|
||||||
|
for turn in history:
|
||||||
|
if turn["role"] == "user":
|
||||||
|
messages.append(HumanMessage(content=turn["content"]))
|
||||||
|
else:
|
||||||
|
messages.append(AIMessage(content=turn["content"]))
|
||||||
|
|
||||||
|
llm = get_agent_llm("setup", temperature=0.4)
|
||||||
|
llm_with_tools = llm.bind_tools(tools)
|
||||||
|
tool_map = {tool_def.name: tool_def for tool_def in tools}
|
||||||
|
|
||||||
|
_lf_ctx = langfuse_context(user_id=user_id or None, session_id=session_id or None)
|
||||||
|
_lf_ctx.__enter__()
|
||||||
|
|
||||||
|
_span_ctx = (
|
||||||
|
lf.start_as_current_observation(
|
||||||
|
as_type="span",
|
||||||
|
name="journey-setup",
|
||||||
|
input=history[-1]["content"] if history else "",
|
||||||
|
)
|
||||||
|
if lf else None
|
||||||
|
)
|
||||||
|
_span = _span_ctx.__enter__() if _span_ctx else None
|
||||||
|
|
||||||
|
try:
|
||||||
|
for step in range(_MAX_TOOL_STEPS):
|
||||||
|
_gen_ctx = (
|
||||||
|
lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="journey-setup-llm",
|
||||||
|
model=model_for_agent("setup"),
|
||||||
|
prompt=langfuse_prompt,
|
||||||
|
input=messages,
|
||||||
|
)
|
||||||
|
if lf else None
|
||||||
|
)
|
||||||
|
_gen = _gen_ctx.__enter__() if _gen_ctx else None
|
||||||
|
response: AIMessage = await llm_with_tools.ainvoke(messages)
|
||||||
|
if _gen_ctx:
|
||||||
|
_gen.update(output=_as_text(response.content), usage_details=extract_usage(response))
|
||||||
|
_gen_ctx.__exit__(None, None, None)
|
||||||
|
|
||||||
|
resp_text = _as_text(response.content)
|
||||||
|
|
||||||
|
# Guard against empty responses (e.g. model returned finish_reason
|
||||||
|
# 'error' which LiteLLM maps to 'stop' with empty content).
|
||||||
|
if not response.tool_calls and not resp_text.strip():
|
||||||
|
logger.warning(
|
||||||
|
"agent_setup: journey LLM returned empty response at step %d — retrying",
|
||||||
|
step,
|
||||||
|
)
|
||||||
|
# Drop the empty AIMessage so we don't pollute history, and retry.
|
||||||
|
continue
|
||||||
|
|
||||||
|
messages.append(response)
|
||||||
|
|
||||||
|
if not response.tool_calls:
|
||||||
|
if _span:
|
||||||
|
_span.update(output=resp_text)
|
||||||
|
return resp_text
|
||||||
|
|
||||||
|
for call in response.tool_calls:
|
||||||
|
call_name = str(call.get("name", ""))
|
||||||
|
call_args = call.get("args", {})
|
||||||
|
logger.info(
|
||||||
|
"agent_setup: journey tool_call name=%s args=%s",
|
||||||
|
call_name,
|
||||||
|
json.dumps(call_args, ensure_ascii=True)[:500],
|
||||||
|
)
|
||||||
|
|
||||||
|
tool_fn = tool_map.get(call_name)
|
||||||
|
if tool_fn is None:
|
||||||
|
tool_output = f"Unknown tool: {call_name}"
|
||||||
|
else:
|
||||||
|
tool_output = await tool_fn.ainvoke(call_args)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"agent_setup: journey tool_result name=%s output=%s",
|
||||||
|
call_name,
|
||||||
|
str(tool_output)[:800],
|
||||||
|
)
|
||||||
|
messages.append(ToolMessage(content=str(tool_output), tool_call_id=call["id"]))
|
||||||
|
|
||||||
|
# Fallback: exceeded max steps.
|
||||||
|
final = await llm.ainvoke(messages)
|
||||||
|
final_text = _as_text(final.content)
|
||||||
|
if _span:
|
||||||
|
_span.update(output=final_text)
|
||||||
|
return final_text or (
|
||||||
|
"Sorry, I had trouble processing the files. "
|
||||||
|
"Could you try again? If the issue persists, the files might be too large for me to analyse."
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
if _span_ctx:
|
||||||
|
_span_ctx.__exit__(None, None, None)
|
||||||
|
_lf_ctx.__exit__(None, None, None)
|
||||||
|
if lf:
|
||||||
|
lf.flush()
|
||||||
|
|
||||||
|
|
||||||
|
# ── Journey handlers (called from device_ws.py) ──────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
async def handle_journey_start(
|
||||||
|
user_id: str,
|
||||||
|
frame: dict[str, Any],
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Handle a ``journey_start`` WS frame.
|
||||||
|
|
||||||
|
Creates a session, runs the setup LLM with directory exploration,
|
||||||
|
and returns the ``journey_reply`` payload.
|
||||||
|
"""
|
||||||
|
agent_type = frame.get("agent_type", "local")
|
||||||
|
directory = frame.get("directory", "")
|
||||||
|
data_types = frame.get("data_types", [])
|
||||||
|
existing_config = frame.get("existing_config")
|
||||||
|
|
||||||
|
# Use the session_id provided by the FE so the reply matches the
|
||||||
|
# listener key; fall back to a generated one if absent.
|
||||||
|
session_id = frame.get("session_id") or str(uuid.uuid4())
|
||||||
|
system_prompt, langfuse_prompt = _build_system_prompt(directory, data_types, existing_config)
|
||||||
|
|
||||||
|
session = JourneySession(
|
||||||
|
session_id=session_id,
|
||||||
|
user_id=user_id,
|
||||||
|
agent_type=agent_type,
|
||||||
|
directory=directory,
|
||||||
|
data_types=data_types,
|
||||||
|
system_prompt=system_prompt,
|
||||||
|
langfuse_prompt=langfuse_prompt,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Seed with an initial user message — some providers require at least one
|
||||||
|
# user/input message to be present.
|
||||||
|
seed_history: list[dict[str, Any]] = [
|
||||||
|
{"role": "user", "content": "Hi, I'm ready to set up my agent. Please explore my directory and ask me your first question."},
|
||||||
|
]
|
||||||
|
ai_reply = await _call_llm_with_tools(
|
||||||
|
system_prompt=system_prompt,
|
||||||
|
history=seed_history,
|
||||||
|
tools=make_directory_tools(directory),
|
||||||
|
user_id=user_id,
|
||||||
|
session_id=session_id,
|
||||||
|
langfuse_prompt=langfuse_prompt,
|
||||||
|
)
|
||||||
|
|
||||||
|
session.history.extend(seed_history)
|
||||||
|
session.history.append({"role": "assistant", "content": ai_reply})
|
||||||
|
_sessions[session_id] = session
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"agent_setup: journey session %s started for user %s (directory=%s)",
|
||||||
|
session_id,
|
||||||
|
user_id,
|
||||||
|
directory,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check if the LLM produced the config on the first turn (unlikely but possible).
|
||||||
|
agent_config = _extract_agent_config(ai_reply)
|
||||||
|
done = agent_config is not None
|
||||||
|
|
||||||
|
display_message = ai_reply
|
||||||
|
if done:
|
||||||
|
display_message = (
|
||||||
|
ai_reply[: ai_reply.index(_CONFIG_START)].strip()
|
||||||
|
or "Here is your agent configuration. You can save it or continue refining."
|
||||||
|
)
|
||||||
|
_sessions.pop(session_id, None)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"type": "journey_reply",
|
||||||
|
"session_id": session_id,
|
||||||
|
"message": display_message,
|
||||||
|
"done": done,
|
||||||
|
"agent_config": agent_config,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
async def handle_journey_message(
|
||||||
|
user_id: str,
|
||||||
|
frame: dict[str, Any],
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Handle a ``journey_message`` WS frame.
|
||||||
|
|
||||||
|
Appends the user message, calls the LLM, and returns the
|
||||||
|
``journey_reply`` payload.
|
||||||
|
"""
|
||||||
|
session_id = frame.get("session_id", "")
|
||||||
|
message = frame.get("message", "")
|
||||||
|
|
||||||
|
session = get_journey_session(session_id, user_id)
|
||||||
|
if session is None:
|
||||||
|
return {
|
||||||
|
"type": "journey_reply",
|
||||||
|
"session_id": session_id,
|
||||||
|
"message": "Journey session not found or expired. Please start a new setup.",
|
||||||
|
"done": True,
|
||||||
|
"agent_config": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Append user turn.
|
||||||
|
session.history.append({"role": "user", "content": message})
|
||||||
|
|
||||||
|
# Call the LLM with tools.
|
||||||
|
session_tools = make_directory_tools(session.directory)
|
||||||
|
ai_reply = await _call_llm_with_tools(
|
||||||
|
system_prompt=session.system_prompt,
|
||||||
|
history=session.history,
|
||||||
|
tools=session_tools,
|
||||||
|
user_id=session.user_id,
|
||||||
|
session_id=session_id,
|
||||||
|
langfuse_prompt=session.langfuse_prompt,
|
||||||
|
)
|
||||||
|
|
||||||
|
session.history.append({"role": "assistant", "content": ai_reply})
|
||||||
|
|
||||||
|
# Check if the LLM produced the final config.
|
||||||
|
agent_config = _extract_agent_config(ai_reply)
|
||||||
|
done = agent_config is not None
|
||||||
|
|
||||||
|
# If the LLM didn't produce a config, nudge it once it hits the hard safety cap.
|
||||||
|
if not done:
|
||||||
|
turns = sum(1 for t in session.history if t["role"] == "user")
|
||||||
|
if turns >= _MAX_TURNS:
|
||||||
|
nudge_content = (
|
||||||
|
"[System: You have enough information. Please generate the final "
|
||||||
|
f"AgentConfig JSON now, wrapped in {_CONFIG_START} / {_CONFIG_END} markers.]"
|
||||||
|
)
|
||||||
|
session.history.append({"role": "user", "content": nudge_content})
|
||||||
|
|
||||||
|
nudge_reply = await _call_llm_with_tools(
|
||||||
|
system_prompt=session.system_prompt,
|
||||||
|
history=session.history,
|
||||||
|
tools=session_tools,
|
||||||
|
user_id=session.user_id,
|
||||||
|
session_id=session_id,
|
||||||
|
langfuse_prompt=session.langfuse_prompt,
|
||||||
|
)
|
||||||
|
session.history.append({"role": "assistant", "content": nudge_reply})
|
||||||
|
|
||||||
|
agent_config = _extract_agent_config(nudge_reply)
|
||||||
|
if agent_config is not None:
|
||||||
|
done = True
|
||||||
|
ai_reply = nudge_reply
|
||||||
|
|
||||||
|
display_message = ai_reply
|
||||||
|
if done:
|
||||||
|
display_message = (
|
||||||
|
ai_reply[: ai_reply.index(_CONFIG_START)].strip()
|
||||||
|
if _CONFIG_START in ai_reply
|
||||||
|
else "Here is your agent configuration. You can save it or continue refining."
|
||||||
|
)
|
||||||
|
_sessions.pop(session_id, None)
|
||||||
|
logger.info("agent_setup: journey session %s completed for user %s", session_id, user_id)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"type": "journey_reply",
|
||||||
|
"session_id": session_id,
|
||||||
|
"message": display_message,
|
||||||
|
"done": done,
|
||||||
|
"agent_config": agent_config,
|
||||||
|
}
|
||||||
257
app/api/routes/agents.py
Normal file
257
app/api/routes/agents.py
Normal file
@@ -0,0 +1,257 @@
|
|||||||
|
"""Agent routes.
|
||||||
|
|
||||||
|
Backend responsibilities are intentionally minimal:
|
||||||
|
GET /agents/catalog — static catalog for UI display
|
||||||
|
POST /agents/can-create — billing eligibility check
|
||||||
|
POST /agents/trigger — trigger a local agent run
|
||||||
|
|
||||||
|
Agent configuration is owned by the Electron app and is not persisted
|
||||||
|
in backend agent-config tables.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends, HTTPException, status
|
||||||
|
from sqlalchemy import func, select
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from app.api.deps import get_current_user
|
||||||
|
from app.billing.tier_manager import FEATURES
|
||||||
|
from app.core.agent_runner import is_agent_running, run_local_agent
|
||||||
|
from app.core.device_manager import device_manager
|
||||||
|
from app.core.note_summarizer import generate_note_summary
|
||||||
|
from app.db import get_session
|
||||||
|
from app.models import AgentRunLog, LocalAgentConfig
|
||||||
|
from app.schemas import (
|
||||||
|
AgentCatalogItem,
|
||||||
|
AgentCreationCheckRequest,
|
||||||
|
AgentCreationCheckResponse,
|
||||||
|
AgentRunLogResponse,
|
||||||
|
AgentTriggerRequest,
|
||||||
|
UserProfile,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/agents", tags=["agents"])
|
||||||
|
|
||||||
|
|
||||||
|
# ── Datetime helpers ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def _dt_ms(dt: datetime) -> int:
|
||||||
|
return int(dt.timestamp() * 1000)
|
||||||
|
|
||||||
|
|
||||||
|
def _dt_ms_opt(dt: datetime | None) -> int | None:
|
||||||
|
return int(dt.timestamp() * 1000) if dt else None
|
||||||
|
|
||||||
|
|
||||||
|
def _to_data_types(values: list[str]) -> list[str]:
|
||||||
|
normalize = {
|
||||||
|
"task": "tasks", "tasks": "tasks",
|
||||||
|
"note": "notes", "notes": "notes",
|
||||||
|
"timeline": "timelines", "timelines": "timelines", "timelineEvents": "timelines",
|
||||||
|
"project": "projects", "projects": "projects",
|
||||||
|
}
|
||||||
|
seen: set[str] = set()
|
||||||
|
result: list[str] = []
|
||||||
|
for v in values:
|
||||||
|
mapped = normalize.get(v)
|
||||||
|
if mapped and mapped not in seen:
|
||||||
|
seen.add(mapped)
|
||||||
|
result.append(mapped)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def _to_run_log_response(log: AgentRunLog) -> AgentRunLogResponse:
|
||||||
|
return AgentRunLogResponse(
|
||||||
|
id=log.id,
|
||||||
|
agent_id=log.agent_id,
|
||||||
|
agent_type=log.agent_type, # type: ignore[arg-type]
|
||||||
|
status=log.status, # type: ignore[arg-type]
|
||||||
|
items_processed=log.items_processed,
|
||||||
|
items_created=log.items_created,
|
||||||
|
errors=log.errors or [],
|
||||||
|
started_at=_dt_ms(log.started_at),
|
||||||
|
completed_at=_dt_ms_opt(log.completed_at),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _enforce_agent_limit(tier: str, current_count: int) -> int:
|
||||||
|
limit: int = FEATURES.get(tier, FEATURES["free"])["batch_active"]
|
||||||
|
if limit != -1 and current_count >= limit:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=status.HTTP_403_FORBIDDEN,
|
||||||
|
detail=f"Agent limit ({limit}) reached for your tier. Upgrade to create more.",
|
||||||
|
)
|
||||||
|
return limit
|
||||||
|
|
||||||
|
|
||||||
|
async def _enforce_run_frequency(
|
||||||
|
tier: str,
|
||||||
|
user_id: str,
|
||||||
|
db: AsyncSession,
|
||||||
|
) -> None:
|
||||||
|
"""Raise HTTP 402 if the user has exceeded their daily batch run limit."""
|
||||||
|
limit: int = FEATURES.get(tier, FEATURES["free"])["batch_runs_per_day"]
|
||||||
|
if limit == -1:
|
||||||
|
return # unlimited
|
||||||
|
|
||||||
|
today_start = datetime.now(timezone.utc).replace(
|
||||||
|
hour=0, minute=0, second=0, microsecond=0
|
||||||
|
)
|
||||||
|
result = await db.execute(
|
||||||
|
select(func.count(AgentRunLog.id)).where(
|
||||||
|
AgentRunLog.user_id == user_id,
|
||||||
|
AgentRunLog.started_at >= today_start,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
runs_today: int = result.scalar_one()
|
||||||
|
|
||||||
|
if runs_today >= limit:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
||||||
|
detail=f"Daily batch run limit ({limit}) reached for your tier. Upgrade for more runs.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Catalog ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
@router.get("/catalog", response_model=list[AgentCatalogItem])
|
||||||
|
async def get_agent_catalog(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
) -> list[AgentCatalogItem]:
|
||||||
|
"""Return the static list of available agent types and their descriptions."""
|
||||||
|
return [
|
||||||
|
AgentCatalogItem(
|
||||||
|
type="local_directory",
|
||||||
|
name="Local Directory Monitor",
|
||||||
|
description="Watches local directories, extracts data from files using AI",
|
||||||
|
),
|
||||||
|
AgentCatalogItem(
|
||||||
|
type="gmail",
|
||||||
|
name="Gmail Connector",
|
||||||
|
description="Scans Gmail inbox, extracts tasks/notes from emails",
|
||||||
|
),
|
||||||
|
AgentCatalogItem(
|
||||||
|
type="teams",
|
||||||
|
name="Microsoft Teams Connector",
|
||||||
|
description="Monitors Teams messages, extracts action items",
|
||||||
|
),
|
||||||
|
AgentCatalogItem(
|
||||||
|
type="outlook",
|
||||||
|
name="Outlook Connector",
|
||||||
|
description="Scans Outlook inbox, extracts tasks/notes",
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/can-create", response_model=AgentCreationCheckResponse)
|
||||||
|
async def can_create_agent(
|
||||||
|
body: AgentCreationCheckRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
) -> AgentCreationCheckResponse:
|
||||||
|
"""Check if the user can create one more agent based on billing tier.
|
||||||
|
|
||||||
|
Since configuration is client-owned, the Electron app sends its current
|
||||||
|
active agent count and the backend applies tier limits.
|
||||||
|
"""
|
||||||
|
limit: int = FEATURES.get(current_user.tier, FEATURES["free"])["batch_active"]
|
||||||
|
allowed = limit == -1 or body.active_agents < limit
|
||||||
|
return AgentCreationCheckResponse(
|
||||||
|
allowed=allowed,
|
||||||
|
tier=current_user.tier,
|
||||||
|
active_agents=body.active_agents,
|
||||||
|
limit=limit,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/trigger", response_model=AgentRunLogResponse, status_code=status.HTTP_202_ACCEPTED)
|
||||||
|
async def trigger_agent_run(
|
||||||
|
body: AgentTriggerRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> AgentRunLogResponse:
|
||||||
|
"""Trigger a local agent run using client-provided configuration."""
|
||||||
|
_enforce_agent_limit(current_user.tier, body.active_agents)
|
||||||
|
await _enforce_run_frequency(current_user.tier, current_user.id, db)
|
||||||
|
|
||||||
|
last_run_dt = (
|
||||||
|
datetime.fromtimestamp(body.last_run_at / 1000, tz=timezone.utc)
|
||||||
|
if body.last_run_at
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
config = LocalAgentConfig(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=current_user.id,
|
||||||
|
device_id=body.device_id,
|
||||||
|
name="Local Directory Monitor",
|
||||||
|
directory_paths=[body.directory],
|
||||||
|
data_types=_to_data_types(body.what_to_extract),
|
||||||
|
prompt_template=body.custom_agent_prompt or "",
|
||||||
|
agent_config=body.agent_config,
|
||||||
|
file_extensions=[],
|
||||||
|
schedule_cron=body.batch_interval,
|
||||||
|
enabled=True,
|
||||||
|
last_run_at=last_run_dt,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Use the FE's stable agent_id if provided, fall back to the ephemeral config id.
|
||||||
|
stable_agent_id = body.agent_id or config.id
|
||||||
|
|
||||||
|
if is_agent_running(stable_agent_id):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=status.HTTP_409_CONFLICT,
|
||||||
|
detail="Agent is already running. Only one run per agent is allowed at a time.",
|
||||||
|
)
|
||||||
|
|
||||||
|
run_log = AgentRunLog(
|
||||||
|
agent_id=stable_agent_id,
|
||||||
|
agent_type="local",
|
||||||
|
user_id=current_user.id,
|
||||||
|
status="running",
|
||||||
|
)
|
||||||
|
db.add(run_log)
|
||||||
|
await db.commit()
|
||||||
|
await db.refresh(run_log)
|
||||||
|
|
||||||
|
run_context = {
|
||||||
|
"type": "agent_batch",
|
||||||
|
"run_id": run_log.id,
|
||||||
|
"agent_id": stable_agent_id,
|
||||||
|
}
|
||||||
|
|
||||||
|
asyncio.create_task(
|
||||||
|
run_local_agent(current_user.id, config, run_log, device_manager, run_context)
|
||||||
|
)
|
||||||
|
|
||||||
|
return _to_run_log_response(run_log)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Note summary endpoint ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class NoteSummarizeRequest(BaseModel):
|
||||||
|
title: str
|
||||||
|
content: str
|
||||||
|
|
||||||
|
|
||||||
|
class NoteSummarizeResponse(BaseModel):
|
||||||
|
summary: str
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/notes/summarize", response_model=NoteSummarizeResponse)
|
||||||
|
async def summarize_note(
|
||||||
|
body: NoteSummarizeRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
) -> NoteSummarizeResponse:
|
||||||
|
"""Generate an AI summary for a note. Used by the Electron backfill on startup."""
|
||||||
|
summary = await generate_note_summary(body.title, body.content)
|
||||||
|
return NoteSummarizeResponse(summary=summary)
|
||||||
@@ -1,33 +1,68 @@
|
|||||||
"""Auth routes: register, login, refresh, me.
|
"""Auth routes: register, login, refresh, me, OAuth social login, onboarding.
|
||||||
|
|
||||||
Users and refresh tokens are persisted in PostgreSQL (users + refresh_tokens
|
Users and refresh tokens are persisted in PostgreSQL (users + refresh_tokens
|
||||||
tables). Passwords are hashed with bcrypt; refresh tokens are stored as
|
tables). Passwords are hashed with bcrypt; refresh tokens are stored as
|
||||||
SHA-256 hashes so plaintext never reaches the DB.
|
SHA-256 hashes so plaintext never reaches the DB.
|
||||||
|
|
||||||
|
OAuth (Google):
|
||||||
|
GET /auth/oauth/{provider}/authorize — returns consent-screen URL + state
|
||||||
|
POST /auth/oauth/{provider}/callback — exchanges code, issues JWT tokens
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import json
|
||||||
import time
|
import time
|
||||||
|
import urllib.parse
|
||||||
import uuid
|
import uuid
|
||||||
from datetime import datetime, timedelta, timezone
|
from datetime import datetime, timedelta, timezone
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
import bcrypt
|
import bcrypt
|
||||||
|
from cryptography.fernet import Fernet
|
||||||
from fastapi import APIRouter, Depends, HTTPException, status
|
from fastapi import APIRouter, Depends, HTTPException, status
|
||||||
|
from fastapi.responses import RedirectResponse
|
||||||
from jose import jwt
|
from jose import jwt
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel, Field
|
||||||
from sqlalchemy import select
|
from sqlalchemy import select
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
from app.api.deps import get_current_user
|
||||||
|
from app.auth.oauth_providers import GoogleOAuthProvider, generate_pkce_pair
|
||||||
from app.config.settings import settings
|
from app.config.settings import settings
|
||||||
|
from app.core.llm import get_llm
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware
|
||||||
from app.db import get_session
|
from app.db import get_session
|
||||||
from app.models import RefreshToken, User
|
from app.models import OAuthAccount, RefreshToken, User
|
||||||
from app.schemas import AuthTokens, UserProfile
|
from app.schemas import AuthTokens, UserProfile
|
||||||
|
|
||||||
router = APIRouter(prefix="/auth", tags=["auth"])
|
router = APIRouter(prefix="/auth", tags=["auth"])
|
||||||
|
|
||||||
|
|
||||||
|
# ── OAuth provider registry ───────────────────────────────────────────
|
||||||
|
|
||||||
|
def _get_google_provider() -> GoogleOAuthProvider:
|
||||||
|
if not settings.GOOGLE_AUTH_CLIENT_ID or not settings.GOOGLE_AUTH_CLIENT_SECRET:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||||
|
"Google login is not configured on this server",
|
||||||
|
)
|
||||||
|
return GoogleOAuthProvider(
|
||||||
|
client_id=settings.GOOGLE_AUTH_CLIENT_ID,
|
||||||
|
client_secret=settings.GOOGLE_AUTH_CLIENT_SECRET,
|
||||||
|
redirect_uri=settings.OAUTH_REDIRECT_URI,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
_PROVIDERS = {"google": _get_google_provider}
|
||||||
|
|
||||||
|
# In-memory state store: state → (code_verifier, expires_at_epoch_s)
|
||||||
|
# Production note: replace with Redis for multi-process deployments.
|
||||||
|
_pending_states: dict[str, tuple[str, float]] = {}
|
||||||
|
_STATE_TTL_SECONDS = 600 # 10 minutes
|
||||||
|
|
||||||
|
|
||||||
# ── Internal helpers ─────────────────────────────────────────────────
|
# ── Internal helpers ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@@ -65,6 +100,8 @@ def _make_access_token(user_id: str, email: str, tier: str) -> tuple[str, int]:
|
|||||||
class _RegisterRequest(BaseModel):
|
class _RegisterRequest(BaseModel):
|
||||||
email: str
|
email: str
|
||||||
password: str
|
password: str
|
||||||
|
name: str | None = None
|
||||||
|
surname: str | None = None
|
||||||
|
|
||||||
|
|
||||||
class _LoginRequest(BaseModel):
|
class _LoginRequest(BaseModel):
|
||||||
@@ -92,8 +129,11 @@ async def register(
|
|||||||
user = User(
|
user = User(
|
||||||
id=str(uuid.uuid4()),
|
id=str(uuid.uuid4()),
|
||||||
email=body.email,
|
email=body.email,
|
||||||
|
name=body.name,
|
||||||
|
surname=body.surname,
|
||||||
password_hash=_hash_password(body.password),
|
password_hash=_hash_password(body.password),
|
||||||
tier="free",
|
tier="free",
|
||||||
|
encryption_key=Fernet.generate_key().decode(),
|
||||||
)
|
)
|
||||||
db.add(user)
|
db.add(user)
|
||||||
await db.flush() # get user.id without committing
|
await db.flush() # get user.id without committing
|
||||||
@@ -191,7 +231,565 @@ async def refresh(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _UpdateProfileRequest(BaseModel):
|
||||||
|
name: str | None = None
|
||||||
|
surname: str | None = None
|
||||||
|
|
||||||
|
|
||||||
@router.get("/me", response_model=UserProfile)
|
@router.get("/me", response_model=UserProfile)
|
||||||
async def me(current_user: UserProfile = Depends(get_current_user)) -> UserProfile:
|
async def me(current_user: UserProfile = Depends(get_current_user)) -> UserProfile:
|
||||||
"""Return the profile for the authenticated user."""
|
"""Return the profile for the authenticated user."""
|
||||||
return current_user
|
return current_user
|
||||||
|
|
||||||
|
|
||||||
|
@router.put("/me", response_model=UserProfile)
|
||||||
|
async def update_profile(
|
||||||
|
body: _UpdateProfileRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> UserProfile:
|
||||||
|
"""Update the authenticated user's name and surname."""
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
|
||||||
|
if body.name is not None:
|
||||||
|
user.name = body.name
|
||||||
|
if body.surname is not None:
|
||||||
|
user.surname = body.surname
|
||||||
|
|
||||||
|
await db.commit()
|
||||||
|
await db.refresh(user)
|
||||||
|
|
||||||
|
return UserProfile(
|
||||||
|
id=user.id,
|
||||||
|
email=user.email,
|
||||||
|
name=user.name,
|
||||||
|
surname=user.surname,
|
||||||
|
avatar_url=user.avatar_url,
|
||||||
|
tier=current_user.tier,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── OAuth helpers ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
async def _issue_refresh_token(user: User, db: AsyncSession) -> tuple[str, AuthTokens]:
|
||||||
|
"""Create a refresh token row and return (plain_token, AuthTokens)."""
|
||||||
|
plain_token = str(uuid.uuid4())
|
||||||
|
expires_at = datetime.now(timezone.utc) + timedelta(
|
||||||
|
days=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS
|
||||||
|
)
|
||||||
|
rt = RefreshToken(
|
||||||
|
user_id=user.id,
|
||||||
|
token_hash=_hash_token(plain_token),
|
||||||
|
expires_at=expires_at,
|
||||||
|
)
|
||||||
|
db.add(rt)
|
||||||
|
access_token, expires_at_ms = _make_access_token(user.id, user.email, user.tier)
|
||||||
|
return plain_token, AuthTokens(
|
||||||
|
access_token=access_token,
|
||||||
|
refresh_token=plain_token,
|
||||||
|
expires_at=expires_at_ms,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── OAuth request/response schemas ───────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class _OAuthAuthorizeResponse(BaseModel):
|
||||||
|
url: str
|
||||||
|
state: str
|
||||||
|
|
||||||
|
|
||||||
|
class _OAuthCallbackRequest(BaseModel):
|
||||||
|
code: str
|
||||||
|
state: str
|
||||||
|
|
||||||
|
|
||||||
|
# ── OAuth routes ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
@router.get(
|
||||||
|
"/oauth/{provider}/web-callback",
|
||||||
|
summary="Web-facing OAuth redirect — bounces to the adiuvai:// deep link",
|
||||||
|
include_in_schema=False,
|
||||||
|
)
|
||||||
|
async def oauth_web_callback(
|
||||||
|
provider: Literal["google"],
|
||||||
|
code: str,
|
||||||
|
state: str,
|
||||||
|
) -> RedirectResponse:
|
||||||
|
"""Google redirects here after user consent.
|
||||||
|
|
||||||
|
This endpoint immediately redirects to the Electron deep-link URI so the
|
||||||
|
desktop app receives the authorization code. It is intentionally simple —
|
||||||
|
no state validation here (the Electron app + backend callback do that).
|
||||||
|
|
||||||
|
Registered in Google Cloud Console as:
|
||||||
|
http://localhost:8000/api/v1/auth/oauth/google/web-callback (dev)
|
||||||
|
https://api.adiuvai.com/api/v1/auth/oauth/google/web-callback (prod)
|
||||||
|
"""
|
||||||
|
params = urllib.parse.urlencode({"code": code, "state": state, "provider": provider})
|
||||||
|
deep_link = f"adiuvai://oauth/callback?{params}"
|
||||||
|
return RedirectResponse(url=deep_link, status_code=302)
|
||||||
|
|
||||||
|
|
||||||
|
@router.get(
|
||||||
|
"/oauth/{provider}/authorize",
|
||||||
|
response_model=_OAuthAuthorizeResponse,
|
||||||
|
summary="Start OAuth flow — returns the provider consent-screen URL",
|
||||||
|
)
|
||||||
|
async def oauth_authorize(
|
||||||
|
provider: Literal["google"],
|
||||||
|
) -> _OAuthAuthorizeResponse:
|
||||||
|
"""Generate a PKCE state + code_challenge and return the authorization URL.
|
||||||
|
|
||||||
|
The client opens this URL in the system browser. After the user grants
|
||||||
|
consent, the provider redirects to the deep-link URI (adiuvai://oauth/callback)
|
||||||
|
with ``code`` and ``state`` query params. The client then calls
|
||||||
|
``POST /auth/oauth/{provider}/callback`` with those values.
|
||||||
|
"""
|
||||||
|
provider_factory = _PROVIDERS.get(provider)
|
||||||
|
if provider_factory is None:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, f"Unknown provider: {provider}")
|
||||||
|
|
||||||
|
oauth_provider = provider_factory()
|
||||||
|
state = str(uuid.uuid4())
|
||||||
|
code_verifier, code_challenge = generate_pkce_pair()
|
||||||
|
|
||||||
|
# Purge expired states to prevent unbounded growth.
|
||||||
|
now = time.time()
|
||||||
|
expired = [s for s, (_, exp) in _pending_states.items() if exp < now]
|
||||||
|
for s in expired:
|
||||||
|
del _pending_states[s]
|
||||||
|
|
||||||
|
_pending_states[state] = (code_verifier, now + _STATE_TTL_SECONDS)
|
||||||
|
|
||||||
|
url = oauth_provider.get_authorization_url(state=state, code_challenge=code_challenge)
|
||||||
|
return _OAuthAuthorizeResponse(url=url, state=state)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post(
|
||||||
|
"/oauth/{provider}/callback",
|
||||||
|
response_model=AuthTokens,
|
||||||
|
summary="Complete OAuth flow — exchange code and issue JWT tokens",
|
||||||
|
)
|
||||||
|
async def oauth_callback(
|
||||||
|
provider: Literal["google"],
|
||||||
|
body: _OAuthCallbackRequest,
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> AuthTokens:
|
||||||
|
"""Validate state, exchange the authorization code, and sign in (or register) the user.
|
||||||
|
|
||||||
|
Resolution order:
|
||||||
|
1. ``oauth_accounts`` row match → existing user, log in.
|
||||||
|
2. Email match + ``email_verified=True`` → link OAuth account to existing user.
|
||||||
|
3. No match → create new user (password_hash=None, avatar from provider).
|
||||||
|
"""
|
||||||
|
provider_factory = _PROVIDERS.get(provider)
|
||||||
|
if provider_factory is None:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, f"Unknown provider: {provider}")
|
||||||
|
|
||||||
|
# Validate state (CSRF protection).
|
||||||
|
now = time.time()
|
||||||
|
entry = _pending_states.pop(body.state, None)
|
||||||
|
if entry is None or entry[1] < now:
|
||||||
|
raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Invalid or expired OAuth state")
|
||||||
|
|
||||||
|
code_verifier, _ = entry
|
||||||
|
|
||||||
|
oauth_provider = provider_factory()
|
||||||
|
|
||||||
|
# Exchange code for tokens.
|
||||||
|
try:
|
||||||
|
token_data = await oauth_provider.exchange_code(
|
||||||
|
code=body.code,
|
||||||
|
code_verifier=code_verifier,
|
||||||
|
redirect_uri=settings.OAUTH_REDIRECT_URI,
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST, "Failed to exchange authorization code"
|
||||||
|
)
|
||||||
|
|
||||||
|
access_token_google = token_data.get("access_token")
|
||||||
|
if not access_token_google:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, "No access token in provider response")
|
||||||
|
|
||||||
|
# Fetch user identity.
|
||||||
|
try:
|
||||||
|
userinfo = await oauth_provider.get_userinfo(access_token_google)
|
||||||
|
except Exception:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, "Failed to fetch user info from provider")
|
||||||
|
|
||||||
|
# ── Resolution order ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
# 1. Existing OAuth link?
|
||||||
|
oauth_result = await db.execute(
|
||||||
|
select(OAuthAccount).where(
|
||||||
|
OAuthAccount.provider == provider,
|
||||||
|
OAuthAccount.provider_user_id == userinfo.provider_user_id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
oauth_account = oauth_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if oauth_account is not None:
|
||||||
|
user_result = await db.execute(select(User).where(User.id == oauth_account.user_id))
|
||||||
|
user = user_result.scalar_one()
|
||||||
|
# Backfill avatar if the user doesn't have one yet.
|
||||||
|
if user.avatar_url is None and userinfo.avatar_url:
|
||||||
|
user.avatar_url = userinfo.avatar_url
|
||||||
|
await db.commit()
|
||||||
|
plain_token, tokens = await _issue_refresh_token(user, db)
|
||||||
|
await db.commit()
|
||||||
|
return tokens
|
||||||
|
|
||||||
|
# 2. Email match with a verified Google email → link accounts.
|
||||||
|
if userinfo.email_verified:
|
||||||
|
email_result = await db.execute(select(User).where(User.email == userinfo.email))
|
||||||
|
existing_user = email_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if existing_user is not None:
|
||||||
|
new_link = OAuthAccount(
|
||||||
|
user_id=existing_user.id,
|
||||||
|
provider=provider,
|
||||||
|
provider_user_id=userinfo.provider_user_id,
|
||||||
|
provider_email=userinfo.email,
|
||||||
|
)
|
||||||
|
db.add(new_link)
|
||||||
|
if existing_user.avatar_url is None and userinfo.avatar_url:
|
||||||
|
existing_user.avatar_url = userinfo.avatar_url
|
||||||
|
plain_token, tokens = await _issue_refresh_token(existing_user, db)
|
||||||
|
await db.commit()
|
||||||
|
return tokens
|
||||||
|
|
||||||
|
# Guard: if the email is already taken but we couldn't auto-link (e.g.
|
||||||
|
# email_verified=False), refuse with 409 instead of hitting a DB constraint.
|
||||||
|
if not userinfo.email_verified:
|
||||||
|
conflict = await db.execute(select(User).where(User.email == userinfo.email))
|
||||||
|
if conflict.scalar_one_or_none() is not None:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_409_CONFLICT,
|
||||||
|
"An account with this email already exists. "
|
||||||
|
"Please sign in with your password.",
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3. New user — social-only account (no password).
|
||||||
|
new_user = User(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
email=userinfo.email,
|
||||||
|
name=userinfo.name,
|
||||||
|
password_hash=None,
|
||||||
|
avatar_url=userinfo.avatar_url,
|
||||||
|
tier="free",
|
||||||
|
encryption_key=Fernet.generate_key().decode(),
|
||||||
|
)
|
||||||
|
db.add(new_user)
|
||||||
|
await db.flush() # populate new_user.id
|
||||||
|
|
||||||
|
new_oauth = OAuthAccount(
|
||||||
|
user_id=new_user.id,
|
||||||
|
provider=provider,
|
||||||
|
provider_user_id=userinfo.provider_user_id,
|
||||||
|
provider_email=userinfo.email,
|
||||||
|
)
|
||||||
|
db.add(new_oauth)
|
||||||
|
|
||||||
|
plain_token, tokens = await _issue_refresh_token(new_user, db)
|
||||||
|
await db.commit()
|
||||||
|
return tokens
|
||||||
|
|
||||||
|
|
||||||
|
# ── Onboarding helpers ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
async def _build_profile(user_id: str, email: str, db: AsyncSession) -> UserProfile:
|
||||||
|
"""Re-fetch and return a full UserProfile (reuses get_current_user logic)."""
|
||||||
|
|
||||||
|
# We can't call the FastAPI dependency directly, but we can replicate
|
||||||
|
# the core logic inline. Instead, we just re-query the same way.
|
||||||
|
from app.models import Subscription # noqa: PLC0415
|
||||||
|
|
||||||
|
result = await db.execute(
|
||||||
|
select(Subscription.tier).where(Subscription.user_id == user_id)
|
||||||
|
)
|
||||||
|
default_tier = "power" if settings.ENV == "dev" else "free"
|
||||||
|
tier: str = result.scalar_one_or_none() or default_tier
|
||||||
|
|
||||||
|
user_result = await db.execute(
|
||||||
|
select(
|
||||||
|
User.name, User.surname, User.avatar_url, User.onboarding_completed_at,
|
||||||
|
User.password_hash,
|
||||||
|
).where(User.id == user_id)
|
||||||
|
)
|
||||||
|
user_row = user_result.one_or_none()
|
||||||
|
|
||||||
|
onboarding_ms: int | None = None
|
||||||
|
if user_row and user_row.onboarding_completed_at is not None:
|
||||||
|
onboarding_ms = int(user_row.onboarding_completed_at.timestamp() * 1000)
|
||||||
|
|
||||||
|
memory_dict: dict[str, str] = {}
|
||||||
|
try:
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
blocks = await mw.list_core_blocks(user_id)
|
||||||
|
memory_dict = {b["label"]: b["value"] for b in blocks}
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return UserProfile(
|
||||||
|
id=user_id,
|
||||||
|
email=email,
|
||||||
|
name=user_row.name if user_row else None,
|
||||||
|
surname=user_row.surname if user_row else None,
|
||||||
|
avatar_url=user_row.avatar_url if user_row else None,
|
||||||
|
has_password=bool(user_row.password_hash) if user_row else False,
|
||||||
|
tier=tier,
|
||||||
|
onboarding_completed_at=onboarding_ms,
|
||||||
|
memory=memory_dict,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Onboarding routes ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class _UpdateMemoryRequest(BaseModel):
|
||||||
|
memory: dict[str, str] = Field(default_factory=dict)
|
||||||
|
mark_onboarded: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
@router.put("/me/memory", response_model=UserProfile)
|
||||||
|
async def update_memory(
|
||||||
|
body: _UpdateMemoryRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> UserProfile:
|
||||||
|
"""Update core memory key/value pairs and optionally mark onboarding complete."""
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
for key, value in body.memory.items():
|
||||||
|
await mw.update_core(current_user.id, key, value)
|
||||||
|
if body.mark_onboarded:
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
user.onboarding_completed_at = datetime.now(timezone.utc)
|
||||||
|
await db.commit()
|
||||||
|
return await _build_profile(current_user.id, current_user.email, db)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/me/onboarding/reset")
|
||||||
|
async def reset_onboarding(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
):
|
||||||
|
"""Reset onboarding so the wizard runs again on next login."""
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
user.onboarding_completed_at = None
|
||||||
|
await db.commit()
|
||||||
|
return {"status": "reset"}
|
||||||
|
|
||||||
|
|
||||||
|
class _NormalizeRequest(BaseModel):
|
||||||
|
inputs: dict[str, str]
|
||||||
|
|
||||||
|
|
||||||
|
class _NormalizeResponse(BaseModel):
|
||||||
|
normalized: dict[str, str]
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/onboarding/normalize", response_model=_NormalizeResponse)
|
||||||
|
async def normalize_onboarding(
|
||||||
|
body: _NormalizeRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
) -> _NormalizeResponse:
|
||||||
|
"""One-shot LLM normalization for free-text onboarding answers."""
|
||||||
|
if not body.inputs:
|
||||||
|
return _NormalizeResponse(normalized={})
|
||||||
|
try:
|
||||||
|
llm = get_llm(model="gpt-4o-mini", temperature=0)
|
||||||
|
prompt = (
|
||||||
|
"You normalize user onboarding answers into clean, ≤3-word canonical labels.\n"
|
||||||
|
"Return a JSON object with the same keys and normalized values.\n"
|
||||||
|
"Examples: 'i build websites' → 'Web Developer', 'tech-ish stuff' → 'Technology'\n"
|
||||||
|
f"Input: {json.dumps(body.inputs)}"
|
||||||
|
)
|
||||||
|
response = await llm.ainvoke(
|
||||||
|
[
|
||||||
|
{"role": "system", "content": "You normalize user inputs. Return JSON only."},
|
||||||
|
{"role": "user", "content": prompt},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
normalized = json.loads(response.content)
|
||||||
|
return _NormalizeResponse(normalized=normalized)
|
||||||
|
except Exception:
|
||||||
|
# LLM failure must never block onboarding — return inputs unchanged
|
||||||
|
return _NormalizeResponse(normalized=body.inputs)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Password management ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class _ChangePasswordRequest(BaseModel):
|
||||||
|
current_password: str = Field(min_length=1)
|
||||||
|
new_password: str = Field(min_length=8)
|
||||||
|
|
||||||
|
|
||||||
|
@router.put("/me/password", status_code=status.HTTP_200_OK)
|
||||||
|
async def change_password(
|
||||||
|
body: _ChangePasswordRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str, bool]:
|
||||||
|
"""Change the authenticated user's password.
|
||||||
|
|
||||||
|
Requires the current password for verification.
|
||||||
|
Returns 400 for social-only users (no password set).
|
||||||
|
"""
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
|
||||||
|
if user.password_hash is None:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST,
|
||||||
|
"This account uses social login and has no password to change",
|
||||||
|
)
|
||||||
|
|
||||||
|
if not _verify_password(body.current_password, user.password_hash):
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, "Current password is incorrect")
|
||||||
|
|
||||||
|
user.password_hash = _hash_password(body.new_password)
|
||||||
|
await db.commit()
|
||||||
|
return {"ok": True}
|
||||||
|
|
||||||
|
|
||||||
|
# ── OAuth account management ─────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/me/oauth-accounts", response_model=list[dict])
|
||||||
|
async def list_oauth_accounts(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> list[dict]:
|
||||||
|
"""List all OAuth providers linked to the authenticated user."""
|
||||||
|
result = await db.execute(
|
||||||
|
select(OAuthAccount).where(OAuthAccount.user_id == current_user.id)
|
||||||
|
)
|
||||||
|
accounts = result.scalars().all()
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"provider": a.provider,
|
||||||
|
"provider_email": a.provider_email,
|
||||||
|
"created_at": int(a.created_at.timestamp() * 1000),
|
||||||
|
}
|
||||||
|
for a in accounts
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/me/oauth-accounts/{provider}", status_code=status.HTTP_200_OK)
|
||||||
|
async def unlink_oauth_account(
|
||||||
|
provider: str,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str, bool]:
|
||||||
|
"""Unlink an OAuth provider from the authenticated user.
|
||||||
|
|
||||||
|
Refuses if the user has no password and this is their only login method.
|
||||||
|
"""
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
|
||||||
|
oauth_result = await db.execute(
|
||||||
|
select(OAuthAccount).where(
|
||||||
|
OAuthAccount.user_id == current_user.id,
|
||||||
|
OAuthAccount.provider == provider,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
account = oauth_result.scalar_one_or_none()
|
||||||
|
if account is None:
|
||||||
|
raise HTTPException(status.HTTP_404_NOT_FOUND, f"No linked {provider} account found")
|
||||||
|
|
||||||
|
# Safety: don't let users lock themselves out.
|
||||||
|
all_oauth = await db.execute(
|
||||||
|
select(OAuthAccount).where(OAuthAccount.user_id == current_user.id)
|
||||||
|
)
|
||||||
|
oauth_count = len(all_oauth.scalars().all())
|
||||||
|
|
||||||
|
if user.password_hash is None and oauth_count <= 1:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST,
|
||||||
|
"Cannot unlink the only login method. Set a password first.",
|
||||||
|
)
|
||||||
|
|
||||||
|
await db.delete(account)
|
||||||
|
await db.commit()
|
||||||
|
return {"ok": True}
|
||||||
|
|
||||||
|
|
||||||
|
# ── Avatar update ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class _UpdateAvatarRequest(BaseModel):
|
||||||
|
avatar_url: str = Field(min_length=1)
|
||||||
|
|
||||||
|
|
||||||
|
@router.put("/me/avatar", response_model=UserProfile)
|
||||||
|
async def update_avatar(
|
||||||
|
body: _UpdateAvatarRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> UserProfile:
|
||||||
|
"""Update the authenticated user's avatar URL.
|
||||||
|
|
||||||
|
Accepts {"avatar_url": "https://..."} — the client uploads the image
|
||||||
|
to its own storage and passes the resulting URL here.
|
||||||
|
"""
|
||||||
|
if not body.avatar_url.startswith(("https://", "http://", "data:image/")):
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, "Invalid avatar URL")
|
||||||
|
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
user.avatar_url = body.avatar_url
|
||||||
|
await db.commit()
|
||||||
|
|
||||||
|
return await _build_profile(current_user.id, current_user.email, db)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Account deletion ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/me", status_code=status.HTTP_200_OK)
|
||||||
|
async def delete_account(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str, bool]:
|
||||||
|
"""Permanently delete the authenticated user's account.
|
||||||
|
|
||||||
|
Cascades: refresh tokens, OAuth accounts, subscription, and all memory
|
||||||
|
rows are deleted via SQLAlchemy relationship cascades. Stripe subscription
|
||||||
|
is cancelled if active.
|
||||||
|
"""
|
||||||
|
# Cancel Stripe subscription if present.
|
||||||
|
try:
|
||||||
|
from app.billing.stripe_service import stripe_service # noqa: PLC0415
|
||||||
|
await stripe_service.cancel_subscription(current_user.id, db)
|
||||||
|
except HTTPException:
|
||||||
|
pass # No subscription — that's fine
|
||||||
|
|
||||||
|
# Delete all memory rows (core, associative, episodic, proactive).
|
||||||
|
try:
|
||||||
|
from app.models import ( # noqa: PLC0415
|
||||||
|
MemoryAssociative, MemoryCore, MemoryEpisodic, MemoryProactive,
|
||||||
|
)
|
||||||
|
for model in (MemoryCore, MemoryAssociative, MemoryEpisodic, MemoryProactive):
|
||||||
|
await db.execute(
|
||||||
|
model.__table__.delete().where(model.user_id == current_user.id)
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
pass # Non-critical — cascade on User will handle most
|
||||||
|
|
||||||
|
# Delete the user row — cascades handle refresh_tokens, oauth_accounts, subscription.
|
||||||
|
result = await db.execute(select(User).where(User.id == current_user.id))
|
||||||
|
user = result.scalar_one()
|
||||||
|
await db.delete(user)
|
||||||
|
await db.commit()
|
||||||
|
|
||||||
|
return {"ok": True}
|
||||||
|
|||||||
@@ -1,171 +0,0 @@
|
|||||||
"""Backup routes: upload, download, history, and delete E2E-encrypted backups.
|
|
||||||
|
|
||||||
Blobs are stored in S3 via BlobStore. Backup metadata is persisted in the
|
|
||||||
PostgreSQL ``backup_metadata`` table.
|
|
||||||
|
|
||||||
IMPORTANT: GET /history must be declared BEFORE GET / to avoid FastAPI
|
|
||||||
treating "history" as a ``{backup_id}`` path parameter.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import uuid
|
|
||||||
from email.utils import parsedate_to_datetime
|
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
|
|
||||||
from sqlalchemy import func, select
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
|
||||||
from app.billing.tier_manager import tier_manager
|
|
||||||
from app.db import get_session
|
|
||||||
from app.models import BackupMetadata as BackupMetadataModel
|
|
||||||
from app.schemas import BackupMetadata, UserProfile
|
|
||||||
from app.storage.blob_store import BlobStore
|
|
||||||
from app.storage.encryption import reject_if_tampered
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/backup", tags=["backup"])
|
|
||||||
|
|
||||||
_blob_store = BlobStore()
|
|
||||||
|
|
||||||
|
|
||||||
async def _current_backup_bytes(user_id: str, db: AsyncSession) -> int:
|
|
||||||
"""Return total backup bytes stored by *user_id*."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(func.coalesce(func.sum(BackupMetadataModel.size_bytes), 0)).where(
|
|
||||||
BackupMetadataModel.user_id == user_id
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return int(result.scalar_one())
|
|
||||||
|
|
||||||
|
|
||||||
async def _check_backup_quota(
|
|
||||||
user: UserProfile, size_bytes: int, db: AsyncSession
|
|
||||||
) -> None:
|
|
||||||
"""Raise HTTP 402 if the upload would exceed the tier's backup limit."""
|
|
||||||
current = await _current_backup_bytes(user.id, db)
|
|
||||||
tier_manager.enforce_backup_quota(
|
|
||||||
user.tier, current_bytes=current, additional_bytes=size_bytes
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@router.put("")
|
|
||||||
async def upload_backup(
|
|
||||||
request: Request,
|
|
||||||
x_backup_version: int = Header(..., alias="X-Backup-Version"),
|
|
||||||
x_backup_timestamp: int = Header(..., alias="X-Backup-Timestamp"),
|
|
||||||
x_backup_checksum: str = Header(..., alias="X-Backup-Checksum"),
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> dict[str, bool]:
|
|
||||||
"""Upload an E2E-encrypted backup blob.
|
|
||||||
|
|
||||||
Metadata is passed via custom headers; the raw body is the encrypted blob.
|
|
||||||
"""
|
|
||||||
blob = await request.body()
|
|
||||||
reject_if_tampered(blob, x_backup_checksum)
|
|
||||||
await _check_backup_quota(current_user, len(blob), db)
|
|
||||||
|
|
||||||
s3_key = await _blob_store.upload(
|
|
||||||
current_user.id, "backup", str(x_backup_timestamp), blob, x_backup_checksum
|
|
||||||
)
|
|
||||||
|
|
||||||
row = BackupMetadataModel(
|
|
||||||
id=str(uuid.uuid4()),
|
|
||||||
user_id=current_user.id,
|
|
||||||
s3_key=s3_key,
|
|
||||||
version=x_backup_version,
|
|
||||||
timestamp=x_backup_timestamp,
|
|
||||||
checksum=x_backup_checksum,
|
|
||||||
size_bytes=len(blob),
|
|
||||||
)
|
|
||||||
db.add(row)
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
return {"ok": True}
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/history", response_model=list[BackupMetadata])
|
|
||||||
async def backup_history(
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> list[BackupMetadata]:
|
|
||||||
"""Return backup metadata records for the authenticated user (no blob bytes)."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(BackupMetadataModel)
|
|
||||||
.where(BackupMetadataModel.user_id == current_user.id)
|
|
||||||
.order_by(BackupMetadataModel.timestamp.desc())
|
|
||||||
)
|
|
||||||
rows = result.scalars().all()
|
|
||||||
return [
|
|
||||||
BackupMetadata(
|
|
||||||
version=r.version,
|
|
||||||
timestamp=r.timestamp,
|
|
||||||
checksum=r.checksum,
|
|
||||||
chunk_count=1,
|
|
||||||
)
|
|
||||||
for r in rows
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("")
|
|
||||||
async def download_backup(
|
|
||||||
request: Request,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> Response:
|
|
||||||
"""Download the latest backup blob. Supports ``If-Modified-Since``."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(BackupMetadataModel)
|
|
||||||
.where(BackupMetadataModel.user_id == current_user.id)
|
|
||||||
.order_by(BackupMetadataModel.timestamp.desc())
|
|
||||||
.limit(1)
|
|
||||||
)
|
|
||||||
latest = result.scalar_one_or_none()
|
|
||||||
if latest is None:
|
|
||||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="No backup found")
|
|
||||||
|
|
||||||
ims_header = request.headers.get("If-Modified-Since")
|
|
||||||
if ims_header:
|
|
||||||
try:
|
|
||||||
ims_dt = parsedate_to_datetime(ims_header)
|
|
||||||
ims_ms = int(ims_dt.timestamp() * 1000)
|
|
||||||
if latest.timestamp <= ims_ms:
|
|
||||||
return Response(status_code=status.HTTP_304_NOT_MODIFIED)
|
|
||||||
except Exception:
|
|
||||||
pass # malformed header — ignore and serve the blob
|
|
||||||
|
|
||||||
blob = await _blob_store.download(current_user.id, latest.s3_key)
|
|
||||||
return Response(
|
|
||||||
content=blob,
|
|
||||||
media_type="application/octet-stream",
|
|
||||||
headers={
|
|
||||||
"X-Backup-Version": str(latest.version),
|
|
||||||
"X-Backup-Timestamp": str(latest.timestamp),
|
|
||||||
"X-Checksum": latest.checksum,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@router.delete("/{backup_id}", response_model=dict)
|
|
||||||
async def delete_backup(
|
|
||||||
backup_id: str,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> dict[str, bool]:
|
|
||||||
"""Delete a specific backup by ID."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(BackupMetadataModel).where(
|
|
||||||
BackupMetadataModel.id == backup_id,
|
|
||||||
BackupMetadataModel.user_id == current_user.id,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
target = result.scalar_one_or_none()
|
|
||||||
if target is None:
|
|
||||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Backup not found")
|
|
||||||
|
|
||||||
await _blob_store.delete(current_user.id, target.s3_key)
|
|
||||||
await db.delete(target)
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
return {"ok": True}
|
|
||||||
@@ -9,7 +9,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, Header, Request, status
|
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
@@ -83,3 +83,50 @@ async def cancel_subscription(
|
|||||||
"""Cancel the active subscription."""
|
"""Cancel the active subscription."""
|
||||||
await stripe_service.cancel_subscription(current_user.id, db)
|
await stripe_service.cancel_subscription(current_user.id, db)
|
||||||
return {"ok": True}
|
return {"ok": True}
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/invoices", response_model=list[dict])
|
||||||
|
async def list_invoices(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Return billing history (invoices) from Stripe.
|
||||||
|
|
||||||
|
Returns an empty list when Stripe is not configured.
|
||||||
|
"""
|
||||||
|
invoices = await stripe_service.list_invoices(current_user.id, db)
|
||||||
|
return invoices
|
||||||
|
|
||||||
|
|
||||||
|
# ── Quota check ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
from app.billing.quota import check_folder_quota, QuotaExceeded # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
class QuotaCheckRequest(BaseModel):
|
||||||
|
feature: str
|
||||||
|
estimated_files: int
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/quota/check")
|
||||||
|
async def quota_check(
|
||||||
|
payload: QuotaCheckRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict:
|
||||||
|
"""Pre-flight folder quota check. 402 if tier limits would be exceeded."""
|
||||||
|
if payload.feature != "folder_index":
|
||||||
|
raise HTTPException(status_code=400, detail="Unknown feature")
|
||||||
|
try:
|
||||||
|
await check_folder_quota(
|
||||||
|
user_id=current_user.id,
|
||||||
|
tier=current_user.tier,
|
||||||
|
estimated_files=payload.estimated_files,
|
||||||
|
db=db,
|
||||||
|
)
|
||||||
|
except QuotaExceeded as exc:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=402,
|
||||||
|
detail={"reason": exc.reason, "message": str(exc)},
|
||||||
|
)
|
||||||
|
return {"ok": True}
|
||||||
|
|||||||
@@ -1,22 +1,40 @@
|
|||||||
"""Chat routes: POST /chat and WebSocket /chat/stream."""
|
"""Chat routes: POST /chat (REST fallback) and POST /chat/embed (text → vector).
|
||||||
|
|
||||||
|
WebSocket chat is handled by the unified device WS endpoint (/api/v1/ws/device).
|
||||||
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import uuid
|
||||||
import json
|
from typing import Literal
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
|
from fastapi import APIRouter, Depends, HTTPException
|
||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
from jose import JWTError, jwt
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
from app.api.deps import get_current_user
|
||||||
from app.config.settings import settings
|
from app.core.brief_agent import run_home_brief, run_project_brief
|
||||||
from app.core.orchestrator import orchestrate, orchestrate_stream
|
from app.core.deep_agent import run_home
|
||||||
|
from app.core.llm import embed
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware
|
||||||
|
from app.db import async_session
|
||||||
from app.schemas import ChatRequest, UserProfile
|
from app.schemas import ChatRequest, UserProfile
|
||||||
|
|
||||||
router = APIRouter(prefix="/chat", tags=["chat"])
|
router = APIRouter(prefix="/chat", tags=["chat"])
|
||||||
|
|
||||||
_HEARTBEAT_INTERVAL = 30 # seconds
|
|
||||||
|
# ── Embed helpers ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class _EmbedRequest(BaseModel):
|
||||||
|
text: str
|
||||||
|
|
||||||
|
|
||||||
|
class _EmbedResponse(BaseModel):
|
||||||
|
vector: list[float]
|
||||||
|
|
||||||
|
|
||||||
|
# ── Endpoints ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@router.post("")
|
@router.post("")
|
||||||
@@ -24,55 +42,75 @@ async def chat(
|
|||||||
body: ChatRequest,
|
body: ChatRequest,
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
) -> JSONResponse:
|
) -> JSONResponse:
|
||||||
"""Route a chat message through the orchestrator.
|
"""REST fallback for home chat when websocket streaming is unavailable."""
|
||||||
|
response = await run_home(
|
||||||
|
user_id=current_user.id,
|
||||||
|
message=body.message,
|
||||||
|
context=body.context.model_dump(),
|
||||||
|
)
|
||||||
|
return JSONResponse(content={"response": response})
|
||||||
|
|
||||||
Returns ``ChatResponse`` for ``execution_mode='direct'``,
|
|
||||||
or ``ExecutionPlan`` for ``execution_mode='plan'``.
|
class _BriefRequest(BaseModel):
|
||||||
|
mode: Literal["home", "project"]
|
||||||
|
project_id: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class _BriefResponse(BaseModel):
|
||||||
|
response: str
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/brief", response_model=_BriefResponse)
|
||||||
|
async def brief(
|
||||||
|
body: _BriefRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
) -> _BriefResponse:
|
||||||
|
"""REST fallback for brief when the device WebSocket is not ready."""
|
||||||
|
if body.mode == "project":
|
||||||
|
if not body.project_id:
|
||||||
|
raise HTTPException(status_code=422, detail="project_id required for project mode")
|
||||||
|
try:
|
||||||
|
uuid.UUID(body.project_id)
|
||||||
|
except ValueError:
|
||||||
|
raise HTTPException(status_code=422, detail="project_id must be a valid UUID")
|
||||||
|
|
||||||
|
request_id = str(uuid.uuid4())
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
memory_context = await memory.enrich_context(
|
||||||
|
current_user.id,
|
||||||
|
"",
|
||||||
|
trace_id=request_id,
|
||||||
|
session_id=request_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
context: dict = {
|
||||||
|
"_debug": {"request_id": request_id, "user_id": current_user.id},
|
||||||
|
**memory_context,
|
||||||
|
}
|
||||||
|
|
||||||
|
chunks: list[str] = []
|
||||||
|
if body.mode == "project":
|
||||||
|
stream = run_project_brief(current_user.id, body.project_id, context) # type: ignore[arg-type]
|
||||||
|
else:
|
||||||
|
stream = run_home_brief(current_user.id, context)
|
||||||
|
|
||||||
|
async for event_type, data in stream:
|
||||||
|
if event_type == "token" and data:
|
||||||
|
chunks.append(str(data))
|
||||||
|
|
||||||
|
return _BriefResponse(response="".join(chunks))
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/embed", response_model=_EmbedResponse)
|
||||||
|
async def embed_text(
|
||||||
|
body: _EmbedRequest,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
) -> _EmbedResponse:
|
||||||
|
"""Generate a 1536-dim embedding vector for the given text.
|
||||||
|
|
||||||
|
Uses ``text-embedding-3-small`` via OpenAI. Auth required (JWT).
|
||||||
|
Used by Electron (vectordb.ts) for local note search.
|
||||||
"""
|
"""
|
||||||
result = await orchestrate(body)
|
vector = await embed(body.text)
|
||||||
return JSONResponse(content=result.model_dump())
|
return _EmbedResponse(vector=vector)
|
||||||
|
|
||||||
|
|
||||||
@router.websocket("/stream")
|
|
||||||
async def chat_stream(websocket: WebSocket) -> None:
|
|
||||||
"""Streaming chat via WebSocket.
|
|
||||||
|
|
||||||
Auth: ``?token=<jwt>`` query param (Bearer not possible during WS handshake).
|
|
||||||
|
|
||||||
Protocol:
|
|
||||||
1. Client sends ``ChatRequest`` as the first JSON text frame.
|
|
||||||
2. Server streams response text chunks.
|
|
||||||
3. Final frame: JSON ``{"done": true, "response": "...", "actions": [...]}``.
|
|
||||||
4. Server pings every 30 s to keep the connection alive.
|
|
||||||
"""
|
|
||||||
# Authenticate before accepting the connection
|
|
||||||
token = websocket.query_params.get("token", "")
|
|
||||||
try:
|
|
||||||
payload = jwt.decode(token, settings.JWT_SECRET, algorithms=[settings.JWT_ALGORITHM])
|
|
||||||
user_id: str | None = payload.get("sub")
|
|
||||||
if not user_id:
|
|
||||||
raise JWTError("missing sub")
|
|
||||||
except JWTError:
|
|
||||||
await websocket.close(code=1008) # 1008 = Policy Violation
|
|
||||||
return
|
|
||||||
|
|
||||||
await websocket.accept()
|
|
||||||
|
|
||||||
try:
|
|
||||||
raw = await websocket.receive_text()
|
|
||||||
body = ChatRequest.model_validate_json(raw)
|
|
||||||
|
|
||||||
async def _heartbeat() -> None:
|
|
||||||
while True:
|
|
||||||
await asyncio.sleep(_HEARTBEAT_INTERVAL)
|
|
||||||
await websocket.send_text(json.dumps({"ping": True}))
|
|
||||||
|
|
||||||
heartbeat_task = asyncio.create_task(_heartbeat())
|
|
||||||
try:
|
|
||||||
async for chunk in orchestrate_stream(body):
|
|
||||||
await websocket.send_text(chunk)
|
|
||||||
finally:
|
|
||||||
heartbeat_task.cancel()
|
|
||||||
|
|
||||||
except WebSocketDisconnect:
|
|
||||||
pass
|
|
||||||
|
|||||||
792
app/api/routes/device_ws.py
Normal file
792
app/api/routes/device_ws.py
Normal file
@@ -0,0 +1,792 @@
|
|||||||
|
"""Device WebSocket endpoint.
|
||||||
|
|
||||||
|
Persistent connection from Electron devices to the backend.
|
||||||
|
|
||||||
|
WS /api/v1/ws/device?token=<jwt>
|
||||||
|
|
||||||
|
Auth: JWT passed as ``?token=`` query parameter (Bearer header is not
|
||||||
|
available during the WebSocket handshake).
|
||||||
|
|
||||||
|
Protocol:
|
||||||
|
1. Client connects → JWT validated → connection accepted.
|
||||||
|
2. Client sends ``device_hello`` frame: ``{ type, device_id, agent_ids }``.
|
||||||
|
3. Backend registers the connection in ``DeviceConnectionManager``.
|
||||||
|
4. Session enters message dispatch loop + heartbeat.
|
||||||
|
|
||||||
|
Incoming frame dispatch:
|
||||||
|
- ``tool_result`` → resolves a pending tool-call Future.
|
||||||
|
- ``journey_start`` → starts a guided setup journey session.
|
||||||
|
- ``journey_message`` → continues a journey conversation.
|
||||||
|
- ``pong`` → heartbeat acknowledgement (updates last-seen).
|
||||||
|
- unknown types → logged, ignored.
|
||||||
|
|
||||||
|
Outgoing heartbeat: ``{ "type": "ping" }`` every 30 s.
|
||||||
|
|
||||||
|
On disconnect:
|
||||||
|
- Unregisters from DeviceConnectionManager.
|
||||||
|
- Marks all in-progress AgentRunLog rows for this user as ``error``
|
||||||
|
with message "device disconnected".
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||||||
|
from jose import JWTError, jwt
|
||||||
|
from sqlalchemy import update
|
||||||
|
|
||||||
|
from app.api.routes.agent_setup import handle_journey_message, handle_journey_start
|
||||||
|
from app.config.settings import settings
|
||||||
|
from app.core.agent_runner import trigger_pending_runs
|
||||||
|
from app.core.brief_agent import run_home_brief, run_project_brief
|
||||||
|
from app.core.deep_agent import run_floating_stream, run_home_stream, run_task_brief_research_stream
|
||||||
|
from app.core.output_formatter import extract_canvas_block
|
||||||
|
from app.core.device_manager import device_manager
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware
|
||||||
|
from app.core.output_formatter import StreamFormatter
|
||||||
|
from app.core.ws_context import clear_client_executor, set_client_executor
|
||||||
|
from app.db import async_session
|
||||||
|
from app.models import AgentRunLog
|
||||||
|
from app.schemas import WsFrameType, WsStreamEnd
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/ws", tags=["device-ws"])
|
||||||
|
|
||||||
|
# ── v7 folder index session state ─────────────────────────────────────
|
||||||
|
# Keyed by sessionId; value: { user_id, project_id, processed, total, cancelled }
|
||||||
|
_index_sessions: dict[str, dict] = {}
|
||||||
|
|
||||||
|
_HEARTBEAT_INTERVAL = 30 # seconds
|
||||||
|
_PONG_TIMEOUT = 10 # seconds — grace window after a ping
|
||||||
|
|
||||||
|
|
||||||
|
@router.websocket("/device")
|
||||||
|
async def device_ws(websocket: WebSocket) -> None:
|
||||||
|
"""Persistent WebSocket endpoint for Electron device connections.
|
||||||
|
|
||||||
|
Authentication is via ``?token=<jwt>`` query parameter.
|
||||||
|
"""
|
||||||
|
# ── 1. Authenticate before accepting ─────────────────────────────
|
||||||
|
token = websocket.query_params.get("token", "")
|
||||||
|
try:
|
||||||
|
payload = jwt.decode(
|
||||||
|
token, settings.JWT_SECRET, algorithms=[settings.JWT_ALGORITHM]
|
||||||
|
)
|
||||||
|
user_id: str | None = payload.get("sub")
|
||||||
|
if not user_id:
|
||||||
|
raise JWTError("missing sub")
|
||||||
|
except JWTError:
|
||||||
|
await websocket.close(code=1008) # Policy Violation
|
||||||
|
return
|
||||||
|
|
||||||
|
await websocket.accept()
|
||||||
|
|
||||||
|
# ── 2. Await device_hello frame ───────────────────────────────────
|
||||||
|
try:
|
||||||
|
raw = await asyncio.wait_for(websocket.receive_text(), timeout=15.0)
|
||||||
|
except (asyncio.TimeoutError, WebSocketDisconnect):
|
||||||
|
await websocket.close(code=1008)
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
hello = json.loads(raw)
|
||||||
|
if hello.get("type") != WsFrameType.device_hello:
|
||||||
|
raise ValueError("expected device_hello as first frame")
|
||||||
|
device_id: str = hello["device_id"]
|
||||||
|
agent_ids: list[str] = hello.get("agent_ids", [])
|
||||||
|
except (KeyError, ValueError, json.JSONDecodeError) as exc:
|
||||||
|
logger.warning("device_ws: invalid device_hello from user=%s: %s", user_id, exc)
|
||||||
|
await websocket.close(code=1008)
|
||||||
|
return
|
||||||
|
|
||||||
|
# ── 3. Register connection ────────────────────────────────────────
|
||||||
|
device_manager.register(user_id, device_id, websocket)
|
||||||
|
logger.info(
|
||||||
|
"device_ws: connected user=%s device=%s agents=%s",
|
||||||
|
user_id,
|
||||||
|
device_id,
|
||||||
|
agent_ids,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Trigger any overdue agent runs now that the device is connected.
|
||||||
|
asyncio.create_task(trigger_pending_runs(user_id, device_id, device_manager))
|
||||||
|
|
||||||
|
# ── 4. Concurrent message loop + heartbeat ────────────────────────
|
||||||
|
try:
|
||||||
|
await asyncio.gather(
|
||||||
|
_message_loop(websocket, user_id),
|
||||||
|
_heartbeat_loop(websocket),
|
||||||
|
)
|
||||||
|
except WebSocketDisconnect:
|
||||||
|
pass
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("device_ws: unhandled exception user=%s: %s", user_id, exc)
|
||||||
|
finally:
|
||||||
|
device_manager.unregister(user_id)
|
||||||
|
logger.info("device_ws: disconnected user=%s device=%s", user_id, device_id)
|
||||||
|
await _mark_runs_disconnected(user_id)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Message dispatch loop ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def _message_loop(websocket: WebSocket, user_id: str) -> None:
|
||||||
|
"""Receive frames from Electron and dispatch to the appropriate handler."""
|
||||||
|
async for raw in websocket.iter_text():
|
||||||
|
try:
|
||||||
|
frame: dict = json.loads(raw)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
logger.warning("device_ws: invalid JSON from user=%s", user_id)
|
||||||
|
continue
|
||||||
|
|
||||||
|
frame_type = frame.get("type")
|
||||||
|
|
||||||
|
if frame_type == WsFrameType.tool_result:
|
||||||
|
call_id = frame.get("id")
|
||||||
|
if call_id:
|
||||||
|
device_manager.resolve_pending_call(user_id, call_id, frame)
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
"device_ws: tool_result missing id from user=%s", user_id
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.home_request:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_home_request(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.floating_request:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_floating_request(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.brief_request:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_brief_request(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.task_brief_request:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_task_brief_request(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.journey_start:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_journey_start(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.journey_message:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_journey_message(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.index_session_start:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_index_session_start(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.index_file_batch:
|
||||||
|
asyncio.create_task(
|
||||||
|
_handle_index_file_batch(websocket, user_id, frame)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif frame_type == WsFrameType.index_session_cancel:
|
||||||
|
await _handle_index_session_cancel(websocket, frame)
|
||||||
|
|
||||||
|
elif frame_type == "pong":
|
||||||
|
# Heartbeat ack — nothing to do, connection is alive.
|
||||||
|
pass
|
||||||
|
|
||||||
|
else:
|
||||||
|
logger.debug(
|
||||||
|
"device_ws: unknown frame type %r from user=%s", frame_type, user_id
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── v3 Chat Handlers ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def _make_ws_executor(websocket: WebSocket, user_id: str):
|
||||||
|
"""Return a callback that sends tool_call frames and awaits tool_result."""
|
||||||
|
async def _executor(payload: dict) -> dict:
|
||||||
|
payload["type"] = WsFrameType.tool_call
|
||||||
|
await websocket.send_text(json.dumps(payload))
|
||||||
|
future = device_manager.create_pending_call(user_id, payload["id"])
|
||||||
|
return await future
|
||||||
|
return _executor
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_home_request(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Handle a home_request frame — streams HomeFormatter output back on the socket."""
|
||||||
|
request_id = frame.get("request_id") or str(uuid4())
|
||||||
|
message: str = frame.get("message", "")
|
||||||
|
session_id: str = frame.get("session_id") or str(uuid4())
|
||||||
|
project_id: str | None = frame.get("project_id") or frame.get("projectId") or None
|
||||||
|
logger.info(
|
||||||
|
"device_ws: home_request_start user=%s req=%s session=%s project=%s msg=%s",
|
||||||
|
user_id,
|
||||||
|
request_id,
|
||||||
|
session_id,
|
||||||
|
project_id,
|
||||||
|
message[:200],
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Memory: enrich context before LLM call ────────────────────────
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
memory_context = await memory.enrich_context(
|
||||||
|
user_id,
|
||||||
|
message,
|
||||||
|
trace_id=request_id,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
context: dict = {
|
||||||
|
"conversation_history": frame.get("conversation_history", []),
|
||||||
|
"_debug": {"request_id": request_id, "session_id": session_id, "user_id": user_id},
|
||||||
|
"format_prefs": frame.get("format_prefs"),
|
||||||
|
**memory_context,
|
||||||
|
}
|
||||||
|
|
||||||
|
executor = await _make_ws_executor(websocket, user_id)
|
||||||
|
set_client_executor(executor)
|
||||||
|
response_chunks: list[str] = []
|
||||||
|
try:
|
||||||
|
event_stream = run_home_stream(user_id, message, context, project_id=project_id)
|
||||||
|
formatter = StreamFormatter(request_id=request_id)
|
||||||
|
async for ws_frame in formatter.format(event_stream):
|
||||||
|
await websocket.send_text(ws_frame.model_dump_json())
|
||||||
|
# Collect text chunks to build the full response for episode storage
|
||||||
|
if ws_frame.type == "stream_text": # type: ignore[union-attr]
|
||||||
|
response_chunks.append(ws_frame.chunk) # type: ignore[union-attr]
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error(
|
||||||
|
"device_ws: home_request failed user=%s req=%s: %s",
|
||||||
|
user_id, request_id, exc,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
clear_client_executor()
|
||||||
|
|
||||||
|
# ── Memory: store episode after response ──────────────────────────
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
await memory.store_episode(
|
||||||
|
user_id, session_id, message, "".join(response_chunks), trace_id=request_id
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"device_ws: home_request_end user=%s req=%s session=%s response_chars=%d",
|
||||||
|
user_id,
|
||||||
|
request_id,
|
||||||
|
session_id,
|
||||||
|
len("".join(response_chunks)),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_floating_request(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Handle a floating_request frame — streams FloatingFormatter output back on the socket."""
|
||||||
|
request_id = frame.get("request_id") or str(uuid4())
|
||||||
|
message: str = frame.get("message", "")
|
||||||
|
session_id: str = frame.get("session_id") or str(uuid4())
|
||||||
|
scope: dict = frame.get("scope", {})
|
||||||
|
logger.info(
|
||||||
|
"device_ws: floating_request_start user=%s req=%s session=%s scope=%s msg=%s",
|
||||||
|
user_id,
|
||||||
|
request_id,
|
||||||
|
session_id,
|
||||||
|
json.dumps(scope, ensure_ascii=True)[:200],
|
||||||
|
message[:200],
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Memory: enrich context before LLM call ────────────────────────
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
memory_context = await memory.enrich_context(
|
||||||
|
user_id,
|
||||||
|
message,
|
||||||
|
trace_id=request_id,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
context: dict = {
|
||||||
|
"conversation_history": frame.get("conversation_history", []),
|
||||||
|
"scope": scope,
|
||||||
|
"_debug": {"request_id": request_id, "session_id": session_id, "user_id": user_id},
|
||||||
|
"format_prefs": frame.get("format_prefs"),
|
||||||
|
**memory_context,
|
||||||
|
}
|
||||||
|
|
||||||
|
executor = await _make_ws_executor(websocket, user_id)
|
||||||
|
set_client_executor(executor)
|
||||||
|
response_chunks: list[str] = []
|
||||||
|
try:
|
||||||
|
event_stream = run_floating_stream(user_id, message, context)
|
||||||
|
formatter = StreamFormatter(request_id=request_id)
|
||||||
|
async for ws_frame in formatter.format(event_stream):
|
||||||
|
await websocket.send_text(ws_frame.model_dump_json())
|
||||||
|
if ws_frame.type == "stream_text": # type: ignore[union-attr]
|
||||||
|
response_chunks.append(ws_frame.chunk) # type: ignore[union-attr]
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error(
|
||||||
|
"device_ws: floating_request failed user=%s req=%s: %s",
|
||||||
|
user_id, request_id, exc,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
clear_client_executor()
|
||||||
|
|
||||||
|
# ── Memory: store episode after response ──────────────────────────
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
await memory.store_episode(
|
||||||
|
user_id, session_id, message, "".join(response_chunks), trace_id=request_id
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"device_ws: floating_request_end user=%s req=%s session=%s response_chars=%d",
|
||||||
|
user_id,
|
||||||
|
request_id,
|
||||||
|
session_id,
|
||||||
|
len("".join(response_chunks)),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_brief_request(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Handle a brief_request frame — streams plain-text brief back on the socket.
|
||||||
|
|
||||||
|
No episode storage — briefs are not conversations.
|
||||||
|
"""
|
||||||
|
import uuid as _uuid
|
||||||
|
|
||||||
|
request_id = frame.get("request_id") or str(uuid4())
|
||||||
|
session_id = frame.get("session_id") or str(uuid4())
|
||||||
|
mode: str = frame.get("mode", "home")
|
||||||
|
project_id: str | None = frame.get("project_id")
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"device_ws: brief_request_start user=%s req=%s mode=%s project_id=%s",
|
||||||
|
user_id, request_id, mode, project_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate project_id for project mode before touching LLM.
|
||||||
|
if mode == "project":
|
||||||
|
try:
|
||||||
|
if not project_id:
|
||||||
|
raise ValueError("project_id required for project mode")
|
||||||
|
_uuid.UUID(project_id)
|
||||||
|
except (ValueError, AttributeError) as exc:
|
||||||
|
logger.warning(
|
||||||
|
"device_ws: brief_request invalid project_id user=%s req=%s: %s",
|
||||||
|
user_id, request_id, exc,
|
||||||
|
)
|
||||||
|
await websocket.send_text(
|
||||||
|
WsStreamEnd(request_id=request_id, error=str(exc)).model_dump_json()
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Enrich context with memory (no user message — use empty string as probe).
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
memory_context = await memory.enrich_context(
|
||||||
|
user_id,
|
||||||
|
"",
|
||||||
|
trace_id=request_id,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
context: dict = {
|
||||||
|
"_debug": {"request_id": request_id, "session_id": session_id, "user_id": user_id},
|
||||||
|
"format_prefs": frame.get("format_prefs"),
|
||||||
|
**memory_context,
|
||||||
|
}
|
||||||
|
|
||||||
|
executor = await _make_ws_executor(websocket, user_id)
|
||||||
|
set_client_executor(executor)
|
||||||
|
try:
|
||||||
|
if mode == "project":
|
||||||
|
event_stream = run_project_brief(user_id, project_id, context) # type: ignore[arg-type]
|
||||||
|
else:
|
||||||
|
event_stream = run_home_brief(user_id, context)
|
||||||
|
|
||||||
|
formatter = StreamFormatter(request_id=request_id)
|
||||||
|
async for ws_frame in formatter.format(event_stream):
|
||||||
|
await websocket.send_text(ws_frame.model_dump_json())
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error(
|
||||||
|
"device_ws: brief_request failed user=%s req=%s: %s",
|
||||||
|
user_id, request_id, exc,
|
||||||
|
)
|
||||||
|
await websocket.send_text(
|
||||||
|
WsStreamEnd(request_id=request_id, error=str(exc)).model_dump_json()
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
clear_client_executor()
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"device_ws: brief_request_end user=%s req=%s mode=%s",
|
||||||
|
user_id, request_id, mode,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── v6 Task Brief Handler ────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_task_brief_request(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Handle a task_brief_request frame — Stage-1 executive assistant deep research.
|
||||||
|
|
||||||
|
Streams the briefing markdown back to the client.
|
||||||
|
On stream_end, emits a ``canvas_draft`` mutation if the agent produced one.
|
||||||
|
"""
|
||||||
|
request_id = frame.get("request_id") or str(uuid4())
|
||||||
|
session_id = frame.get("session_id") or str(uuid4())
|
||||||
|
task_id: str = frame.get("task_id") or frame.get("taskId") or ""
|
||||||
|
project_id: str | None = frame.get("project_id") or frame.get("projectId") or None
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"device_ws: task_brief_request_start user=%s req=%s task=%s project=%s [cache_miss]",
|
||||||
|
user_id, request_id, task_id, project_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not task_id:
|
||||||
|
await websocket.send_text(
|
||||||
|
WsStreamEnd(request_id=request_id, error="task_id is required").model_dump_json()
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
async with async_session() as db:
|
||||||
|
memory = MemoryMiddleware(db)
|
||||||
|
memory_context = await memory.enrich_context(
|
||||||
|
user_id,
|
||||||
|
f"task brief: {task_id}",
|
||||||
|
trace_id=request_id,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
context: dict = {
|
||||||
|
"_debug": {"request_id": request_id, "session_id": session_id, "user_id": user_id},
|
||||||
|
"format_prefs": frame.get("format_prefs"),
|
||||||
|
**memory_context,
|
||||||
|
}
|
||||||
|
|
||||||
|
executor = await _make_ws_executor(websocket, user_id)
|
||||||
|
set_client_executor(executor)
|
||||||
|
response_chunks: list[str] = []
|
||||||
|
|
||||||
|
try:
|
||||||
|
event_stream = run_task_brief_research_stream(user_id, task_id, context, project_id=project_id)
|
||||||
|
formatter = StreamFormatter(request_id=request_id)
|
||||||
|
async for ws_frame in formatter.format(event_stream):
|
||||||
|
if ws_frame.type == "stream_text": # type: ignore[union-attr]
|
||||||
|
response_chunks.append(ws_frame.chunk) # type: ignore[union-attr]
|
||||||
|
await websocket.send_text(ws_frame.model_dump_json())
|
||||||
|
elif ws_frame.type == "stream_start":
|
||||||
|
await websocket.send_text(ws_frame.model_dump_json())
|
||||||
|
# stream_end is emitted below with mutations — skip formatter's version
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error(
|
||||||
|
"device_ws: task_brief_request failed user=%s req=%s task=%s: %s",
|
||||||
|
user_id, request_id, task_id, exc,
|
||||||
|
)
|
||||||
|
await websocket.send_text(
|
||||||
|
WsStreamEnd(request_id=request_id, error=str(exc)).model_dump_json()
|
||||||
|
)
|
||||||
|
return
|
||||||
|
finally:
|
||||||
|
clear_client_executor()
|
||||||
|
|
||||||
|
# Extract canvas block then emit stream_end with optional mutations.
|
||||||
|
full_response = "".join(response_chunks)
|
||||||
|
_visible, canvas_content, canvas_kind = extract_canvas_block(full_response)
|
||||||
|
|
||||||
|
mutations: list[dict] = []
|
||||||
|
if canvas_content:
|
||||||
|
mutations.append({
|
||||||
|
"type": "canvas_draft",
|
||||||
|
"content": canvas_content,
|
||||||
|
"kind": canvas_kind,
|
||||||
|
})
|
||||||
|
|
||||||
|
await websocket.send_text(
|
||||||
|
WsStreamEnd(request_id=request_id, mutations=mutations or None).model_dump_json()
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"device_ws: task_brief_request_end user=%s req=%s task=%s response_chars=%d canvas=%s",
|
||||||
|
user_id, request_id, task_id, len(full_response), canvas_kind or "none",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── v4 Journey Handlers ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_journey_start(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Handle a journey_start frame — explores directory and sends first question."""
|
||||||
|
executor = await _make_ws_executor(websocket, user_id)
|
||||||
|
set_client_executor(executor)
|
||||||
|
try:
|
||||||
|
reply = await handle_journey_start(user_id, frame)
|
||||||
|
await websocket.send_text(json.dumps(reply))
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error(
|
||||||
|
"device_ws: journey_start failed user=%s: %s", user_id, exc
|
||||||
|
)
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": "journey_reply",
|
||||||
|
"session_id": frame.get("session_id", ""),
|
||||||
|
"message": f"Failed to start journey: {exc}",
|
||||||
|
"done": True,
|
||||||
|
"prompt_template": None,
|
||||||
|
}))
|
||||||
|
finally:
|
||||||
|
clear_client_executor()
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_journey_message(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Handle a journey_message frame — continues the journey conversation."""
|
||||||
|
executor = await _make_ws_executor(websocket, user_id)
|
||||||
|
set_client_executor(executor)
|
||||||
|
try:
|
||||||
|
reply = await handle_journey_message(user_id, frame)
|
||||||
|
await websocket.send_text(json.dumps(reply))
|
||||||
|
except Exception as exc:
|
||||||
|
session_id = frame.get("session_id", "")
|
||||||
|
logger.error(
|
||||||
|
"device_ws: journey_message failed user=%s session=%s: %s",
|
||||||
|
user_id, session_id, exc,
|
||||||
|
)
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": "journey_reply",
|
||||||
|
"session_id": session_id,
|
||||||
|
"message": f"Journey error: {exc}",
|
||||||
|
"done": True,
|
||||||
|
"prompt_template": None,
|
||||||
|
}))
|
||||||
|
finally:
|
||||||
|
clear_client_executor()
|
||||||
|
|
||||||
|
|
||||||
|
# ── v7 Folder Index Handlers ──────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_index_session_start(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Register a new folder index session. No response sent — client is declaring intent."""
|
||||||
|
session_id: str = frame.get("sessionId") or frame.get("session_id") or ""
|
||||||
|
project_id: str | None = frame.get("projectId") or frame.get("project_id")
|
||||||
|
total: int = int(frame.get("totalFiles") or frame.get("total_files") or 0)
|
||||||
|
|
||||||
|
if not session_id:
|
||||||
|
logger.warning("device_ws: index_session_start missing sessionId user=%s", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
_index_sessions[session_id] = {
|
||||||
|
"user_id": user_id,
|
||||||
|
"project_id": project_id,
|
||||||
|
"processed": 0,
|
||||||
|
"total": total,
|
||||||
|
"cancelled": False,
|
||||||
|
}
|
||||||
|
logger.info(
|
||||||
|
"device_ws: index_session_start user=%s session=%s project=%s total=%d",
|
||||||
|
user_id, session_id, project_id, total,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_index_session_cancel(
|
||||||
|
websocket: WebSocket,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Mark a session as cancelled and emit index_session_done(cancelled)."""
|
||||||
|
session_id: str = frame.get("sessionId") or frame.get("session_id") or ""
|
||||||
|
session = _index_sessions.get(session_id)
|
||||||
|
if session:
|
||||||
|
session["cancelled"] = True
|
||||||
|
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": WsFrameType.index_session_done,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"status": "cancelled",
|
||||||
|
}))
|
||||||
|
_index_sessions.pop(session_id, None)
|
||||||
|
logger.info("device_ws: index_session_cancel session=%s", session_id)
|
||||||
|
|
||||||
|
|
||||||
|
async def _handle_index_file_batch(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str,
|
||||||
|
frame: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Process a batch of files for an index session, streaming results back."""
|
||||||
|
# Lazy imports to avoid heavy load at module startup.
|
||||||
|
from app.core.folder_indexer import ( # noqa: PLC0415
|
||||||
|
summarize_image,
|
||||||
|
summarize_pdf,
|
||||||
|
summarize_docx,
|
||||||
|
summarize_text,
|
||||||
|
)
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
from app.billing.quota import add_token_usage # noqa: PLC0415
|
||||||
|
|
||||||
|
session_id: str = frame.get("sessionId") or frame.get("session_id") or ""
|
||||||
|
files: list[dict] = frame.get("files", [])
|
||||||
|
|
||||||
|
session = _index_sessions.get(session_id)
|
||||||
|
if not session or session.get("cancelled"):
|
||||||
|
return
|
||||||
|
|
||||||
|
async with async_session() as db:
|
||||||
|
tier = await tier_manager.get_tier(user_id, db)
|
||||||
|
raw_cap = tier_manager.get_feature_value(tier, "folder_monthly_tokens")
|
||||||
|
cap: int | None = None if raw_cap == -1 else raw_cap
|
||||||
|
|
||||||
|
for file_info in files:
|
||||||
|
if session.get("cancelled"):
|
||||||
|
return
|
||||||
|
|
||||||
|
# Electron's toSnakeCase converts payload keys, so accept both forms.
|
||||||
|
rel_path: str = file_info.get("relPath") or file_info.get("rel_path") or ""
|
||||||
|
kind: str = file_info.get("kind") or "text"
|
||||||
|
content: str = file_info.get("content") or ""
|
||||||
|
ext: str = file_info.get("ext") or ""
|
||||||
|
mime: str = file_info.get("mime") or "application/octet-stream"
|
||||||
|
name: str = rel_path.split("/")[-1] or rel_path
|
||||||
|
|
||||||
|
try:
|
||||||
|
if kind == "image":
|
||||||
|
res = await summarize_image(image_b64=content, mime=mime)
|
||||||
|
elif kind == "pdf":
|
||||||
|
res = await summarize_pdf(pdf_b64=content, name=name)
|
||||||
|
elif kind == "docx":
|
||||||
|
res = await summarize_docx(docx_b64=content, name=name)
|
||||||
|
else:
|
||||||
|
res = await summarize_text(content=content, ext=ext, name=name)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"device_ws: index_file_batch summarize failed session=%s path=%s: %s",
|
||||||
|
session_id, rel_path, exc,
|
||||||
|
)
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": WsFrameType.index_file_result,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"relPath": rel_path,
|
||||||
|
"summary": None,
|
||||||
|
"tokensUsed": 0,
|
||||||
|
"error": str(exc),
|
||||||
|
}))
|
||||||
|
session["processed"] += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Account for token usage and check cap.
|
||||||
|
usage = await add_token_usage(
|
||||||
|
user_id=user_id,
|
||||||
|
feature="folder_index",
|
||||||
|
tokens=res.tokens_used,
|
||||||
|
db=db,
|
||||||
|
cap=cap,
|
||||||
|
)
|
||||||
|
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": WsFrameType.index_file_result,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"relPath": rel_path,
|
||||||
|
"summary": res.summary,
|
||||||
|
"tokensUsed": res.tokens_used,
|
||||||
|
}))
|
||||||
|
session["processed"] += 1
|
||||||
|
|
||||||
|
if usage.exhausted:
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": WsFrameType.index_session_done,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"status": "quota_exceeded",
|
||||||
|
}))
|
||||||
|
_index_sessions.pop(session_id, None)
|
||||||
|
logger.info(
|
||||||
|
"device_ws: index_session quota_exceeded user=%s session=%s",
|
||||||
|
user_id, session_id,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# After processing the batch, emit progress.
|
||||||
|
processed = session["processed"]
|
||||||
|
total = session["total"]
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": WsFrameType.index_session_progress,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"processed": processed,
|
||||||
|
"total": total,
|
||||||
|
}))
|
||||||
|
|
||||||
|
if processed >= total:
|
||||||
|
await websocket.send_text(json.dumps({
|
||||||
|
"type": WsFrameType.index_session_done,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"status": "completed",
|
||||||
|
}))
|
||||||
|
_index_sessions.pop(session_id, None)
|
||||||
|
logger.info(
|
||||||
|
"device_ws: index_session_done completed user=%s session=%s processed=%d",
|
||||||
|
user_id, session_id, processed,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Heartbeat ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def _heartbeat_loop(websocket: WebSocket) -> None:
|
||||||
|
"""Send a ping frame every 30 s to keep the connection alive."""
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(_HEARTBEAT_INTERVAL)
|
||||||
|
await websocket.send_text(json.dumps({"type": "ping"}))
|
||||||
|
|
||||||
|
|
||||||
|
# ── Disconnect cleanup ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def _mark_runs_disconnected(user_id: str) -> None:
|
||||||
|
"""Mark all in-progress AgentRunLog rows as 'error' for this user."""
|
||||||
|
try:
|
||||||
|
async with async_session() as db:
|
||||||
|
await db.execute(
|
||||||
|
update(AgentRunLog)
|
||||||
|
.where(
|
||||||
|
AgentRunLog.user_id == user_id,
|
||||||
|
AgentRunLog.status == "running",
|
||||||
|
)
|
||||||
|
.values(
|
||||||
|
status="error",
|
||||||
|
errors=["device disconnected"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await db.commit()
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error(
|
||||||
|
"device_ws: failed to mark runs as disconnected for user=%s: %s",
|
||||||
|
user_id,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
225
app/api/routes/memory.py
Normal file
225
app/api/routes/memory.py
Normal file
@@ -0,0 +1,225 @@
|
|||||||
|
"""Memory management routes — view/edit/delete user memory tiers.
|
||||||
|
|
||||||
|
All routes require authentication. Data is always user-scoped.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Annotated
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends, Header, HTTPException, status
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from sqlalchemy import delete, select
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.api.deps import get_current_user
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware
|
||||||
|
from app.db import get_session
|
||||||
|
from app.models import (
|
||||||
|
ExtractionQueue,
|
||||||
|
MemoryAssociative,
|
||||||
|
MemoryCore,
|
||||||
|
MemoryEpisodic,
|
||||||
|
MemoryProactive,
|
||||||
|
MemoryRelation,
|
||||||
|
)
|
||||||
|
from app.schemas import UserProfile
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/memory", tags=["memory"])
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_ALLOWED_PREDICATES = {
|
||||||
|
"works_at",
|
||||||
|
"reports_to",
|
||||||
|
"stakeholder_of",
|
||||||
|
"last_contacted_on",
|
||||||
|
"owes_followup",
|
||||||
|
"manages",
|
||||||
|
"collaborates_with",
|
||||||
|
"owns",
|
||||||
|
"member_of",
|
||||||
|
"custom",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# ── Response schemas ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
class RelationOut(BaseModel):
|
||||||
|
id: str
|
||||||
|
subject_label: str
|
||||||
|
subject_type: str
|
||||||
|
predicate: str
|
||||||
|
object_label: str
|
||||||
|
object_type: str
|
||||||
|
confidence: float
|
||||||
|
last_confirmed_at: int | None = None # epoch ms
|
||||||
|
|
||||||
|
|
||||||
|
class RelationPatch(BaseModel):
|
||||||
|
subject_label: str | None = None
|
||||||
|
object_label: str | None = None
|
||||||
|
predicate: str | None = None
|
||||||
|
confidence: float | None = Field(None, ge=0.0, le=1.0)
|
||||||
|
|
||||||
|
|
||||||
|
class CoreAddBody(BaseModel):
|
||||||
|
key: str = Field(..., min_length=1, max_length=255)
|
||||||
|
value: str = Field(..., min_length=1)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Helpers ──────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def _relation_to_out(row: MemoryRelation) -> RelationOut:
|
||||||
|
last_ms: int | None = None
|
||||||
|
if row.last_confirmed_at is not None:
|
||||||
|
last_ms = int(row.last_confirmed_at.timestamp() * 1000)
|
||||||
|
return RelationOut(
|
||||||
|
id=row.id,
|
||||||
|
subject_label=row.subject_label,
|
||||||
|
subject_type=row.subject_type,
|
||||||
|
predicate=row.predicate,
|
||||||
|
object_label=row.object_label,
|
||||||
|
object_type=row.object_type,
|
||||||
|
confidence=row.confidence,
|
||||||
|
last_confirmed_at=last_ms,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Routes ───────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
@router.get("/core", response_model=dict[str, str])
|
||||||
|
async def get_core_memory(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str, str]:
|
||||||
|
"""Return all core memory k/v pairs (plaintext) for the current user."""
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
blocks = await mw.list_core_blocks(current_user.id)
|
||||||
|
return {b["label"]: b["value"] for b in blocks}
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/core/{key}", status_code=status.HTTP_204_NO_CONTENT)
|
||||||
|
async def delete_core_key(
|
||||||
|
key: str,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> None:
|
||||||
|
"""Delete a single core memory key (GDPR Art. 17)."""
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
deleted = await mw.delete_core(current_user.id, key)
|
||||||
|
if not deleted:
|
||||||
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Key not found")
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/core", status_code=status.HTTP_201_CREATED, response_model=dict[str, str])
|
||||||
|
async def add_core_key(
|
||||||
|
body: CoreAddBody,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str, str]:
|
||||||
|
"""Add or overwrite a core memory key/value pair."""
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
await mw.update_core(current_user.id, body.key, body.value)
|
||||||
|
return {body.key: body.value}
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/relational", response_model=list[RelationOut])
|
||||||
|
async def get_relational_memory(
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> list[RelationOut]:
|
||||||
|
"""Return all relational memory rows for the current user."""
|
||||||
|
mw = MemoryMiddleware(db)
|
||||||
|
rows = await mw.query_relations(current_user.id, limit=200)
|
||||||
|
return [_relation_to_out(r) for r in rows]
|
||||||
|
|
||||||
|
|
||||||
|
@router.patch("/relational/{relation_id}", response_model=RelationOut)
|
||||||
|
async def patch_relation(
|
||||||
|
relation_id: str,
|
||||||
|
body: RelationPatch,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> RelationOut:
|
||||||
|
"""Edit a relation row's labels, predicate, or confidence."""
|
||||||
|
if body.predicate is not None and body.predicate not in _ALLOWED_PREDICATES:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||||
|
detail=f"predicate must be one of: {sorted(_ALLOWED_PREDICATES)}",
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await db.execute(
|
||||||
|
select(MemoryRelation).where(
|
||||||
|
MemoryRelation.id == relation_id,
|
||||||
|
MemoryRelation.user_id == current_user.id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
row = result.scalar_one_or_none()
|
||||||
|
if row is None:
|
||||||
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Relation not found")
|
||||||
|
|
||||||
|
if body.subject_label is not None:
|
||||||
|
row.subject_label = body.subject_label
|
||||||
|
if body.object_label is not None:
|
||||||
|
row.object_label = body.object_label
|
||||||
|
if body.predicate is not None:
|
||||||
|
row.predicate = body.predicate
|
||||||
|
if body.confidence is not None:
|
||||||
|
row.confidence = body.confidence
|
||||||
|
row.last_confirmed_at = datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
await db.commit()
|
||||||
|
await db.refresh(row)
|
||||||
|
logger.info("memory: patch_relation user=%s relation=%s", current_user.id, relation_id)
|
||||||
|
return _relation_to_out(row)
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/relational/{relation_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||||||
|
async def delete_relation(
|
||||||
|
relation_id: str,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> None:
|
||||||
|
"""Hard-delete a relation row (GDPR Art. 17)."""
|
||||||
|
result = await db.execute(
|
||||||
|
select(MemoryRelation).where(
|
||||||
|
MemoryRelation.id == relation_id,
|
||||||
|
MemoryRelation.user_id == current_user.id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
row = result.scalar_one_or_none()
|
||||||
|
if row is None:
|
||||||
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Relation not found")
|
||||||
|
await db.delete(row)
|
||||||
|
await db.commit()
|
||||||
|
logger.info("memory: delete_relation user=%s relation=%s", current_user.id, relation_id)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/forget-all", status_code=status.HTTP_204_NO_CONTENT)
|
||||||
|
async def forget_all(
|
||||||
|
x_confirm: Annotated[str | None, Header(alias="X-Confirm")] = None,
|
||||||
|
current_user: UserProfile = Depends(get_current_user),
|
||||||
|
db: AsyncSession = Depends(get_session),
|
||||||
|
) -> None:
|
||||||
|
"""Wipe all memory tiers for the current user (GDPR Art. 17).
|
||||||
|
|
||||||
|
Requires ``X-Confirm: true`` header. Does NOT delete the user account.
|
||||||
|
"""
|
||||||
|
if x_confirm != "true":
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=status.HTTP_400_BAD_REQUEST,
|
||||||
|
detail="Missing or invalid X-Confirm header. Send X-Confirm: true to confirm.",
|
||||||
|
)
|
||||||
|
|
||||||
|
uid = current_user.id
|
||||||
|
await db.execute(delete(MemoryCore).where(MemoryCore.user_id == uid))
|
||||||
|
await db.execute(delete(MemoryAssociative).where(MemoryAssociative.user_id == uid))
|
||||||
|
await db.execute(delete(MemoryEpisodic).where(MemoryEpisodic.user_id == uid))
|
||||||
|
await db.execute(delete(MemoryProactive).where(MemoryProactive.user_id == uid))
|
||||||
|
await db.execute(delete(MemoryRelation).where(MemoryRelation.user_id == uid))
|
||||||
|
await db.execute(delete(ExtractionQueue).where(ExtractionQueue.user_id == uid))
|
||||||
|
await db.commit()
|
||||||
|
logger.warning("memory: forget_all GDPR wipe user=%s", uid)
|
||||||
@@ -1,37 +0,0 @@
|
|||||||
"""Plans routes: GET /plans/playbook and GET /plans/playbook/{plan_id}."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException, status
|
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
|
||||||
from app.core.execution_plan import plan_cache
|
|
||||||
from app.schemas import ExecutionPlan, UserProfile
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/plans", tags=["plans"])
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/playbook", response_model=list[ExecutionPlan])
|
|
||||||
async def list_playbooks(
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
) -> list[ExecutionPlan]:
|
|
||||||
"""Return all cached execution plan playbooks for the authenticated user.
|
|
||||||
|
|
||||||
TODO(Step11): filter by tier — power+ plans gated behind batch_builder feature.
|
|
||||||
"""
|
|
||||||
return plan_cache.get_all_playbooks()
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/playbook/{plan_id}", response_model=ExecutionPlan)
|
|
||||||
async def get_playbook(
|
|
||||||
plan_id: str,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
) -> ExecutionPlan:
|
|
||||||
"""Return a specific execution plan playbook by ID."""
|
|
||||||
plan = plan_cache.get_plan(plan_id)
|
|
||||||
if plan is None:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_404_NOT_FOUND,
|
|
||||||
detail=f"Plan not found: {plan_id}",
|
|
||||||
)
|
|
||||||
return plan
|
|
||||||
@@ -1,148 +0,0 @@
|
|||||||
"""Plugins routes: browse and install plugins from the marketplace.
|
|
||||||
|
|
||||||
Backed by ``PluginRegistry`` and ``RevenueShare`` service classes that
|
|
||||||
persist data in the PostgreSQL ``plugins`` and ``revenue_events`` tables.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from typing import Any, Literal
|
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException, Query, status
|
|
||||||
from pydantic import BaseModel
|
|
||||||
from sqlalchemy import select
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
|
||||||
from app.db import get_session
|
|
||||||
from app.marketplace.plugin_registry import registry
|
|
||||||
from app.marketplace.revenue_share import revenue_share
|
|
||||||
from app.models import PluginInstallation, PluginReview as PluginReviewModel
|
|
||||||
from app.schemas import PluginInstallRequest, PluginListResponse, PluginManifest, UserProfile
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/plugins", tags=["plugins"])
|
|
||||||
|
|
||||||
|
|
||||||
# ── Tier gate ─────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def _require_plugin_tier(user: UserProfile) -> None:
|
|
||||||
"""Raise HTTP 403 for users below Power tier."""
|
|
||||||
if user.tier not in ("power", "team"):
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_403_FORBIDDEN,
|
|
||||||
detail="Plugin marketplace requires Power tier or above",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Local detail schema ────────────────────────────────────────────────
|
|
||||||
|
|
||||||
class _PluginDetail(BaseModel):
|
|
||||||
plugin: PluginManifest
|
|
||||||
install_count: int
|
|
||||||
ratings: list[Any]
|
|
||||||
|
|
||||||
|
|
||||||
# ── Routes ────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
@router.get("", response_model=PluginListResponse)
|
|
||||||
async def list_plugins(
|
|
||||||
category: str | None = Query(default=None),
|
|
||||||
q: str | None = Query(default=None),
|
|
||||||
page: int = Query(default=1, ge=1),
|
|
||||||
sort: Literal["rating", "installs", "newest"] = Query(default="newest"),
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> PluginListResponse:
|
|
||||||
"""Browse the plugin marketplace. Requires Power tier or above."""
|
|
||||||
_require_plugin_tier(current_user)
|
|
||||||
return await registry.list_plugins(db, category=category, query=q, page=page, sort=sort)
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/{plugin_id}", response_model=_PluginDetail)
|
|
||||||
async def get_plugin(
|
|
||||||
plugin_id: str,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> _PluginDetail:
|
|
||||||
"""Get full plugin details including install count. Requires Power tier or above."""
|
|
||||||
_require_plugin_tier(current_user)
|
|
||||||
entry = await registry.get_plugin(db, plugin_id)
|
|
||||||
if entry is None:
|
|
||||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Plugin not found")
|
|
||||||
|
|
||||||
# Fetch review ratings for this plugin
|
|
||||||
review_result = await db.execute(
|
|
||||||
select(PluginReviewModel).where(PluginReviewModel.plugin_id == plugin_id)
|
|
||||||
)
|
|
||||||
reviews = review_result.scalars().all()
|
|
||||||
ratings = [
|
|
||||||
{
|
|
||||||
"reviewer_id": r.reviewer_id,
|
|
||||||
"decision": r.decision,
|
|
||||||
"notes": r.notes,
|
|
||||||
"reviewed_at": int(r.reviewed_at.timestamp() * 1000) if r.reviewed_at else None,
|
|
||||||
}
|
|
||||||
for r in reviews
|
|
||||||
]
|
|
||||||
|
|
||||||
return _PluginDetail(
|
|
||||||
plugin=entry["manifest"],
|
|
||||||
install_count=entry["install_count"],
|
|
||||||
ratings=ratings,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/{plugin_id}/install", response_model=dict)
|
|
||||||
async def install_plugin(
|
|
||||||
plugin_id: str,
|
|
||||||
body: PluginInstallRequest, # noqa: ARG001 — reserved for future fields
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
"""Install a plugin. Triggers Stripe Connect revenue split for paid plugins.
|
|
||||||
|
|
||||||
Requires Power tier or above.
|
|
||||||
"""
|
|
||||||
_require_plugin_tier(current_user)
|
|
||||||
entry = await registry.get_plugin(db, plugin_id)
|
|
||||||
if entry is None:
|
|
||||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Plugin not found")
|
|
||||||
|
|
||||||
# Record the installation in plugin_installations
|
|
||||||
installation = PluginInstallation(
|
|
||||||
plugin_id=plugin_id,
|
|
||||||
user_id=current_user.id,
|
|
||||||
)
|
|
||||||
db.add(installation)
|
|
||||||
await db.flush()
|
|
||||||
|
|
||||||
await revenue_share.record_install(
|
|
||||||
db,
|
|
||||||
plugin_id=plugin_id,
|
|
||||||
user_id=current_user.id,
|
|
||||||
amount_cents=entry["manifest"].price_cents,
|
|
||||||
)
|
|
||||||
|
|
||||||
download_url = f"https://cdn.adiuva.app/plugins/{plugin_id}/package.zip"
|
|
||||||
return {"ok": True, "download_url": download_url}
|
|
||||||
|
|
||||||
|
|
||||||
@router.delete("/{plugin_id}/install", response_model=dict)
|
|
||||||
async def uninstall_plugin(
|
|
||||||
plugin_id: str,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> dict[str, bool]:
|
|
||||||
"""Unregister a plugin installation."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(PluginInstallation).where(
|
|
||||||
PluginInstallation.plugin_id == plugin_id,
|
|
||||||
PluginInstallation.user_id == current_user.id,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
installation = result.scalar_one_or_none()
|
|
||||||
if installation is not None:
|
|
||||||
await db.delete(installation)
|
|
||||||
await db.commit()
|
|
||||||
await registry.record_uninstall(db, plugin_id)
|
|
||||||
return {"ok": True}
|
|
||||||
@@ -1,195 +0,0 @@
|
|||||||
"""Storage routes: CRUD for E2E-encrypted cloud records.
|
|
||||||
|
|
||||||
Blobs are stored in S3 via BlobStore. Record metadata is persisted in the
|
|
||||||
PostgreSQL ``storage_records`` table.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import uuid
|
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
|
|
||||||
from pydantic import BaseModel
|
|
||||||
from sqlalchemy import func, select
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
|
||||||
from app.billing.tier_manager import tier_manager
|
|
||||||
from app.db import get_session
|
|
||||||
from app.models import StorageRecord
|
|
||||||
from app.schemas import StorageRecordCreate, StorageRecordUpdate, UserProfile
|
|
||||||
from app.storage.blob_store import BlobStore
|
|
||||||
from app.storage.encryption import reject_if_tampered
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/storage", tags=["storage"])
|
|
||||||
|
|
||||||
_blob_store = BlobStore()
|
|
||||||
|
|
||||||
|
|
||||||
# ── Local response schemas ─────────────────────────────────────────────
|
|
||||||
|
|
||||||
class _CreateResponse(BaseModel):
|
|
||||||
id: str
|
|
||||||
created_at: int
|
|
||||||
|
|
||||||
|
|
||||||
class _RecordMeta(BaseModel):
|
|
||||||
id: str
|
|
||||||
table: str
|
|
||||||
checksum: str
|
|
||||||
created_at: int
|
|
||||||
updated_at: int
|
|
||||||
|
|
||||||
|
|
||||||
# ── Helpers ────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
async def _current_usage_bytes(user_id: str, db: AsyncSession) -> int:
|
|
||||||
"""Return total bytes stored by *user_id*."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(func.coalesce(func.sum(StorageRecord.size_bytes), 0)).where(
|
|
||||||
StorageRecord.user_id == user_id
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return int(result.scalar_one())
|
|
||||||
|
|
||||||
|
|
||||||
async def _check_quota(user: UserProfile, additional_bytes: int, db: AsyncSession) -> None:
|
|
||||||
"""Raise HTTP 402 if adding *additional_bytes* would exceed the tier limit."""
|
|
||||||
current = await _current_usage_bytes(user.id, db)
|
|
||||||
tier_manager.enforce_quota(user.tier, current_bytes=current, additional_bytes=additional_bytes)
|
|
||||||
|
|
||||||
|
|
||||||
async def _get_record_for_user(
|
|
||||||
record_id: str, user_id: str, db: AsyncSession
|
|
||||||
) -> StorageRecord:
|
|
||||||
"""Look up a record and verify ownership. Returns 404 on mismatch
|
|
||||||
to prevent user enumeration attacks."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(StorageRecord).where(
|
|
||||||
StorageRecord.id == record_id, StorageRecord.user_id == user_id
|
|
||||||
)
|
|
||||||
)
|
|
||||||
record = result.scalar_one_or_none()
|
|
||||||
if record is None:
|
|
||||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Record not found")
|
|
||||||
return record
|
|
||||||
|
|
||||||
|
|
||||||
# ── Routes ─────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
@router.post("/records", response_model=_CreateResponse, status_code=status.HTTP_201_CREATED)
|
|
||||||
async def create_record(
|
|
||||||
body: StorageRecordCreate,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> _CreateResponse:
|
|
||||||
"""Upload a new E2E-encrypted blob. Verifies checksum before storing."""
|
|
||||||
reject_if_tampered(body.blob, body.checksum)
|
|
||||||
await _check_quota(current_user, len(body.blob), db)
|
|
||||||
|
|
||||||
record_id = str(uuid.uuid4())
|
|
||||||
|
|
||||||
s3_key = await _blob_store.upload(
|
|
||||||
current_user.id, body.table, record_id, body.blob, body.checksum
|
|
||||||
)
|
|
||||||
|
|
||||||
record = StorageRecord(
|
|
||||||
id=record_id,
|
|
||||||
user_id=current_user.id,
|
|
||||||
table_name=body.table,
|
|
||||||
s3_key=s3_key,
|
|
||||||
checksum=body.checksum,
|
|
||||||
size_bytes=len(body.blob),
|
|
||||||
)
|
|
||||||
db.add(record)
|
|
||||||
await db.commit()
|
|
||||||
await db.refresh(record)
|
|
||||||
|
|
||||||
created_at_ms = int(record.created_at.timestamp() * 1000)
|
|
||||||
return _CreateResponse(id=record_id, created_at=created_at_ms)
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/records", response_model=list[_RecordMeta])
|
|
||||||
async def list_records(
|
|
||||||
table: str | None = Query(default=None),
|
|
||||||
page: int = Query(default=1, ge=1),
|
|
||||||
limit: int = Query(default=50, ge=1, le=200),
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> list[_RecordMeta]:
|
|
||||||
"""List record metadata for the authenticated user. Blob bytes are never returned."""
|
|
||||||
query = select(StorageRecord).where(StorageRecord.user_id == current_user.id)
|
|
||||||
if table is not None:
|
|
||||||
query = query.where(StorageRecord.table_name == table)
|
|
||||||
query = query.offset((page - 1) * limit).limit(limit)
|
|
||||||
|
|
||||||
result = await db.execute(query)
|
|
||||||
rows = result.scalars().all()
|
|
||||||
|
|
||||||
return [
|
|
||||||
_RecordMeta(
|
|
||||||
id=r.id,
|
|
||||||
table=r.table_name,
|
|
||||||
checksum=r.checksum,
|
|
||||||
created_at=int(r.created_at.timestamp() * 1000),
|
|
||||||
updated_at=int(r.updated_at.timestamp() * 1000),
|
|
||||||
)
|
|
||||||
for r in rows
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/records/{record_id}")
|
|
||||||
async def download_record(
|
|
||||||
record_id: str,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> Response:
|
|
||||||
"""Download an E2E-encrypted blob. Returns raw bytes with ``X-Checksum`` header."""
|
|
||||||
record = await _get_record_for_user(record_id, current_user.id, db)
|
|
||||||
blob = await _blob_store.download(current_user.id, record.s3_key)
|
|
||||||
return Response(
|
|
||||||
content=blob,
|
|
||||||
media_type="application/octet-stream",
|
|
||||||
headers={"X-Checksum": record.checksum},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@router.put("/records/{record_id}", response_model=dict)
|
|
||||||
async def update_record(
|
|
||||||
record_id: str,
|
|
||||||
body: StorageRecordUpdate,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> dict[str, bool]:
|
|
||||||
"""Replace the blob for an existing record. Verifies checksum before storing."""
|
|
||||||
record = await _get_record_for_user(record_id, current_user.id, db)
|
|
||||||
reject_if_tampered(body.blob, body.checksum)
|
|
||||||
|
|
||||||
delta = len(body.blob) - record.size_bytes
|
|
||||||
if delta > 0:
|
|
||||||
await _check_quota(current_user, delta, db)
|
|
||||||
|
|
||||||
s3_key = await _blob_store.upload(
|
|
||||||
current_user.id, record.table_name, record_id, body.blob, body.checksum
|
|
||||||
)
|
|
||||||
|
|
||||||
record.s3_key = s3_key
|
|
||||||
record.checksum = body.checksum
|
|
||||||
record.size_bytes = len(body.blob)
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
return {"ok": True}
|
|
||||||
|
|
||||||
|
|
||||||
@router.delete("/records/{record_id}", response_model=dict)
|
|
||||||
async def delete_record(
|
|
||||||
record_id: str,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
db: AsyncSession = Depends(get_session),
|
|
||||||
) -> dict[str, bool]:
|
|
||||||
"""Delete a record and its S3 blob."""
|
|
||||||
record = await _get_record_for_user(record_id, current_user.id, db)
|
|
||||||
await _blob_store.delete(current_user.id, record.s3_key)
|
|
||||||
await db.delete(record)
|
|
||||||
await db.commit()
|
|
||||||
return {"ok": True}
|
|
||||||
@@ -1,79 +0,0 @@
|
|||||||
"""Vectors routes: upsert, search, delete cloud vector store entries, and embed text."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends
|
|
||||||
from pydantic import BaseModel
|
|
||||||
|
|
||||||
from app.api.deps import get_current_user
|
|
||||||
from app.core.llm import embed
|
|
||||||
from app.schemas import (
|
|
||||||
UserProfile,
|
|
||||||
VectorSearchRequest,
|
|
||||||
VectorSearchResponse,
|
|
||||||
VectorUpsertRequest,
|
|
||||||
)
|
|
||||||
from app.storage.encryption import reject_if_tampered
|
|
||||||
from app.storage.vector_store import VectorStore
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/storage", tags=["vectors"])
|
|
||||||
|
|
||||||
_vector_store = VectorStore()
|
|
||||||
|
|
||||||
|
|
||||||
class _VectorDeleteRequest(BaseModel):
|
|
||||||
ids: list[str]
|
|
||||||
|
|
||||||
|
|
||||||
class _EmbedRequest(BaseModel):
|
|
||||||
text: str
|
|
||||||
|
|
||||||
|
|
||||||
class _EmbedResponse(BaseModel):
|
|
||||||
vector: list[float]
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/vectors/upsert", response_model=dict)
|
|
||||||
async def upsert_vectors(
|
|
||||||
body: VectorUpsertRequest,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
) -> dict[str, int]:
|
|
||||||
"""Verify checksums and store encrypted vectors in the user-scoped namespace."""
|
|
||||||
for item in body.vectors:
|
|
||||||
reject_if_tampered(item.blob, item.checksum)
|
|
||||||
await _vector_store.upsert(current_user.id, body.vectors)
|
|
||||||
return {"upserted": len(body.vectors)}
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/vectors/search", response_model=VectorSearchResponse)
|
|
||||||
async def search_vectors(
|
|
||||||
body: VectorSearchRequest,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
) -> VectorSearchResponse:
|
|
||||||
"""Search the user-scoped vector namespace with an encrypted query blob."""
|
|
||||||
results = await _vector_store.search(current_user.id, body.query_blob, body.top_k)
|
|
||||||
return VectorSearchResponse(results=results)
|
|
||||||
|
|
||||||
|
|
||||||
@router.delete("/vectors", response_model=dict)
|
|
||||||
async def delete_vectors(
|
|
||||||
body: _VectorDeleteRequest,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
) -> dict[str, bool]:
|
|
||||||
"""Delete vectors by ID, scoped to the authenticated user."""
|
|
||||||
await _vector_store.delete(current_user.id, body.ids)
|
|
||||||
return {"ok": True}
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/vectors/embed", response_model=_EmbedResponse)
|
|
||||||
async def embed_text(
|
|
||||||
body: _EmbedRequest,
|
|
||||||
current_user: UserProfile = Depends(get_current_user),
|
|
||||||
) -> _EmbedResponse:
|
|
||||||
"""Generate a 1536-dim embedding vector for the given text.
|
|
||||||
|
|
||||||
Uses ``text-embedding-3-small`` via OpenAI. Auth required (JWT).
|
|
||||||
Used by backend tools (note_agent) and Electron (vectordb.ts) alike.
|
|
||||||
"""
|
|
||||||
vector = await embed(body.text)
|
|
||||||
return _EmbedResponse(vector=vector)
|
|
||||||
1
app/auth/__init__.py
Normal file
1
app/auth/__init__.py
Normal file
@@ -0,0 +1 @@
|
|||||||
|
"OAuth provider abstractions and utilities."
|
||||||
135
app/auth/oauth_providers.py
Normal file
135
app/auth/oauth_providers.py
Normal file
@@ -0,0 +1,135 @@
|
|||||||
|
"""OAuth 2.0 + PKCE provider abstractions.
|
||||||
|
|
||||||
|
Each provider implements a three-step flow designed for a desktop (public) client:
|
||||||
|
|
||||||
|
1. get_authorization_url(state, code_challenge) → str
|
||||||
|
Build the provider's consent-screen URL. State and code_challenge are
|
||||||
|
generated server-side; the client opens this URL in the system browser.
|
||||||
|
|
||||||
|
2. exchange_code(code, code_verifier, redirect_uri) → dict
|
||||||
|
Exchange the short-lived authorization code for an access token.
|
||||||
|
The code_verifier proves ownership of the PKCE challenge.
|
||||||
|
|
||||||
|
3. get_userinfo(access_token) → OAuthUserInfo
|
||||||
|
Fetch the canonical user identity from the provider.
|
||||||
|
|
||||||
|
Currently supported providers:
|
||||||
|
- GoogleOAuthProvider (scope: openid email profile)
|
||||||
|
|
||||||
|
Adding a new provider:
|
||||||
|
- Implement the three methods above.
|
||||||
|
- Register in _PROVIDERS inside routes/auth.py.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import hashlib
|
||||||
|
import os
|
||||||
|
import urllib.parse
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
|
||||||
|
# ── Data transfer objects ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class OAuthUserInfo:
|
||||||
|
"""Normalized user identity returned by any provider."""
|
||||||
|
|
||||||
|
provider_user_id: str
|
||||||
|
email: str
|
||||||
|
email_verified: bool
|
||||||
|
avatar_url: str | None
|
||||||
|
name: str | None
|
||||||
|
|
||||||
|
|
||||||
|
# ── PKCE helpers ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def generate_pkce_pair() -> tuple[str, str]:
|
||||||
|
"""Generate a (code_verifier, code_challenge) pair for PKCE S256.
|
||||||
|
|
||||||
|
The code_verifier is a random 32-byte URL-safe base64 string.
|
||||||
|
The code_challenge is SHA-256(code_verifier) base64url-encoded (no padding).
|
||||||
|
"""
|
||||||
|
code_verifier = base64.urlsafe_b64encode(os.urandom(32)).rstrip(b"=").decode()
|
||||||
|
digest = hashlib.sha256(code_verifier.encode()).digest()
|
||||||
|
code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
|
||||||
|
return code_verifier, code_challenge
|
||||||
|
|
||||||
|
|
||||||
|
# ── Google provider ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class GoogleOAuthProvider:
|
||||||
|
"""Google OAuth 2.0 provider (openid email profile scope).
|
||||||
|
|
||||||
|
Uses Google's standard authorization endpoint with PKCE S256.
|
||||||
|
Does NOT use google-auth-oauthlib to keep the flow generic and async.
|
||||||
|
"""
|
||||||
|
|
||||||
|
name = "google"
|
||||||
|
|
||||||
|
_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
|
||||||
|
_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
||||||
|
_USERINFO_URL = "https://www.googleapis.com/oauth2/v3/userinfo"
|
||||||
|
|
||||||
|
def __init__(self, client_id: str, client_secret: str, redirect_uri: str) -> None:
|
||||||
|
self.client_id = client_id
|
||||||
|
self.client_secret = client_secret
|
||||||
|
self.redirect_uri = redirect_uri
|
||||||
|
|
||||||
|
def get_authorization_url(self, state: str, code_challenge: str) -> str:
|
||||||
|
"""Build the Google consent-screen URL."""
|
||||||
|
params = {
|
||||||
|
"client_id": self.client_id,
|
||||||
|
"redirect_uri": self.redirect_uri,
|
||||||
|
"response_type": "code",
|
||||||
|
"scope": "openid email profile",
|
||||||
|
"state": state,
|
||||||
|
"code_challenge": code_challenge,
|
||||||
|
"code_challenge_method": "S256",
|
||||||
|
"access_type": "offline",
|
||||||
|
"prompt": "select_account",
|
||||||
|
}
|
||||||
|
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
|
||||||
|
|
||||||
|
async def exchange_code(
|
||||||
|
self, code: str, code_verifier: str, redirect_uri: str
|
||||||
|
) -> dict:
|
||||||
|
"""Exchange authorization code for an access token."""
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
response = await client.post(
|
||||||
|
self._TOKEN_URL,
|
||||||
|
data={
|
||||||
|
"client_id": self.client_id,
|
||||||
|
"client_secret": self.client_secret,
|
||||||
|
"code": code,
|
||||||
|
"code_verifier": code_verifier,
|
||||||
|
"grant_type": "authorization_code",
|
||||||
|
"redirect_uri": redirect_uri,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
return response.json()
|
||||||
|
|
||||||
|
async def get_userinfo(self, access_token: str) -> OAuthUserInfo:
|
||||||
|
"""Fetch the authenticated user's identity from Google."""
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
response = await client.get(
|
||||||
|
self._USERINFO_URL,
|
||||||
|
headers={"Authorization": f"Bearer {access_token}"},
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
data = response.json()
|
||||||
|
|
||||||
|
return OAuthUserInfo(
|
||||||
|
provider_user_id=data["sub"],
|
||||||
|
email=data["email"],
|
||||||
|
email_verified=data.get("email_verified", False),
|
||||||
|
avatar_url=data.get("picture"),
|
||||||
|
name=data.get("name"),
|
||||||
|
)
|
||||||
139
app/billing/quota.py
Normal file
139
app/billing/quota.py
Normal file
@@ -0,0 +1,139 @@
|
|||||||
|
"""Quota checks and atomic token-usage accounting for folder integration."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
|
from sqlalchemy import select, update
|
||||||
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.billing.tier_manager import TierManager
|
||||||
|
from app.models import MonthlyTokenUsage
|
||||||
|
from app.schemas import BillingTier
|
||||||
|
|
||||||
|
|
||||||
|
class QuotaExceeded(Exception):
|
||||||
|
"""Raised when a folder operation cannot proceed under the user's tier."""
|
||||||
|
|
||||||
|
def __init__(self, reason: str, message: str) -> None:
|
||||||
|
super().__init__(message)
|
||||||
|
self.reason = reason # "max_files" | "monthly_tokens"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class TokenUsageResult:
|
||||||
|
tokens_used: int
|
||||||
|
exhausted: bool
|
||||||
|
|
||||||
|
|
||||||
|
def _current_year_month() -> str:
|
||||||
|
return datetime.now(timezone.utc).strftime("%Y-%m")
|
||||||
|
|
||||||
|
|
||||||
|
_tier_manager = TierManager()
|
||||||
|
|
||||||
|
|
||||||
|
async def check_folder_quota(
|
||||||
|
*,
|
||||||
|
user_id: str,
|
||||||
|
tier: BillingTier,
|
||||||
|
estimated_files: int,
|
||||||
|
db: AsyncSession,
|
||||||
|
) -> None:
|
||||||
|
"""Raise QuotaExceeded if folder_max_files or folder_monthly_tokens
|
||||||
|
would be violated. -1 in either feature means unlimited."""
|
||||||
|
max_files = _tier_manager.get_feature_value(tier, "folder_max_files")
|
||||||
|
if max_files != -1 and estimated_files > max_files:
|
||||||
|
raise QuotaExceeded(
|
||||||
|
"max_files",
|
||||||
|
f"Folder has {estimated_files} files; tier '{tier}' allows max {max_files}.",
|
||||||
|
)
|
||||||
|
|
||||||
|
cap = _tier_manager.get_feature_value(tier, "folder_monthly_tokens")
|
||||||
|
if cap == -1:
|
||||||
|
return
|
||||||
|
ym = _current_year_month()
|
||||||
|
row = (
|
||||||
|
await db.execute(
|
||||||
|
select(MonthlyTokenUsage).where(
|
||||||
|
MonthlyTokenUsage.user_id == user_id,
|
||||||
|
MonthlyTokenUsage.year_month == ym,
|
||||||
|
MonthlyTokenUsage.feature == "folder_index",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
).scalar_one_or_none()
|
||||||
|
used = row.tokens_used if row else 0
|
||||||
|
if used >= cap:
|
||||||
|
raise QuotaExceeded(
|
||||||
|
"monthly_tokens",
|
||||||
|
f"Monthly token budget exhausted ({used}/{cap}); resets next month.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def add_token_usage(
|
||||||
|
*,
|
||||||
|
user_id: str,
|
||||||
|
feature: str,
|
||||||
|
tokens: int,
|
||||||
|
db: AsyncSession,
|
||||||
|
cap: int | None = None,
|
||||||
|
) -> TokenUsageResult:
|
||||||
|
"""Atomically add `tokens` to MonthlyTokenUsage row for (user, current month, feature).
|
||||||
|
|
||||||
|
Uses PostgreSQL ``INSERT … ON CONFLICT DO UPDATE`` when available; falls
|
||||||
|
back to a read-then-write on other engines (e.g. aiosqlite in tests).
|
||||||
|
Returns post-update total and whether cap is exhausted.
|
||||||
|
"""
|
||||||
|
ym = _current_year_month()
|
||||||
|
|
||||||
|
# Detect dialect to choose between native upsert and portable fallback.
|
||||||
|
dialect_name: str = db.bind.dialect.name if db.bind is not None else "" # type: ignore[union-attr]
|
||||||
|
|
||||||
|
if dialect_name == "postgresql":
|
||||||
|
# Native atomic upsert — production path.
|
||||||
|
stmt = (
|
||||||
|
pg_insert(MonthlyTokenUsage)
|
||||||
|
.values(
|
||||||
|
user_id=user_id,
|
||||||
|
year_month=ym,
|
||||||
|
feature=feature,
|
||||||
|
tokens_used=tokens,
|
||||||
|
)
|
||||||
|
.on_conflict_do_update(
|
||||||
|
index_elements=["user_id", "year_month", "feature"],
|
||||||
|
set_={"tokens_used": MonthlyTokenUsage.tokens_used + tokens},
|
||||||
|
)
|
||||||
|
.returning(MonthlyTokenUsage.tokens_used)
|
||||||
|
)
|
||||||
|
used: int = (await db.execute(stmt)).scalar_one()
|
||||||
|
await db.commit()
|
||||||
|
else:
|
||||||
|
# Portable fallback — used in tests (SQLite) and any non-PG engine.
|
||||||
|
row = (
|
||||||
|
await db.execute(
|
||||||
|
select(MonthlyTokenUsage).where(
|
||||||
|
MonthlyTokenUsage.user_id == user_id,
|
||||||
|
MonthlyTokenUsage.year_month == ym,
|
||||||
|
MonthlyTokenUsage.feature == feature,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
).scalar_one_or_none()
|
||||||
|
|
||||||
|
if row is None:
|
||||||
|
row = MonthlyTokenUsage(
|
||||||
|
user_id=user_id,
|
||||||
|
year_month=ym,
|
||||||
|
feature=feature,
|
||||||
|
tokens_used=tokens,
|
||||||
|
)
|
||||||
|
db.add(row)
|
||||||
|
else:
|
||||||
|
row.tokens_used += tokens
|
||||||
|
|
||||||
|
await db.commit()
|
||||||
|
await db.refresh(row)
|
||||||
|
used = row.tokens_used
|
||||||
|
|
||||||
|
exhausted = cap is not None and cap != -1 and used >= cap
|
||||||
|
return TokenUsageResult(tokens_used=used, exhausted=exhausted)
|
||||||
@@ -43,8 +43,8 @@ class StripeService:
|
|||||||
self,
|
self,
|
||||||
user_id: str,
|
user_id: str,
|
||||||
tier: str,
|
tier: str,
|
||||||
success_url: str = "https://app.adiuva.app/billing/success?session_id={CHECKOUT_SESSION_ID}",
|
success_url: str = "https://app.adiuvai.app/billing/success?session_id={CHECKOUT_SESSION_ID}",
|
||||||
cancel_url: str = "https://app.adiuva.app/billing/cancel",
|
cancel_url: str = "https://app.adiuvai.app/billing/cancel",
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Create a Stripe checkout session and return the URL.
|
"""Create a Stripe checkout session and return the URL.
|
||||||
|
|
||||||
@@ -200,6 +200,45 @@ class StripeService:
|
|||||||
sub.status = "canceled"
|
sub.status = "canceled"
|
||||||
await db.commit()
|
await db.commit()
|
||||||
|
|
||||||
|
async def list_invoices(
|
||||||
|
self, user_id: str, db: AsyncSession, limit: int = 24
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Return recent invoices for the user from Stripe.
|
||||||
|
|
||||||
|
Returns an empty list when Stripe is not configured or the user has
|
||||||
|
no ``stripe_customer_id``.
|
||||||
|
"""
|
||||||
|
if not self._configured():
|
||||||
|
return []
|
||||||
|
|
||||||
|
from app.models import User # noqa: PLC0415
|
||||||
|
|
||||||
|
result = await db.execute(
|
||||||
|
select(User.stripe_customer_id).where(User.id == user_id)
|
||||||
|
)
|
||||||
|
customer_id = result.scalar_one_or_none()
|
||||||
|
if not customer_id:
|
||||||
|
return []
|
||||||
|
|
||||||
|
try:
|
||||||
|
s = self._client()
|
||||||
|
invoices = s.Invoice.list(customer=customer_id, limit=limit)
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"id": inv.id,
|
||||||
|
"amount_due": inv.amount_due,
|
||||||
|
"amount_paid": inv.amount_paid,
|
||||||
|
"currency": inv.currency,
|
||||||
|
"status": inv.status,
|
||||||
|
"created": inv.created * 1000, # epoch ms
|
||||||
|
"invoice_url": inv.hosted_invoice_url,
|
||||||
|
"invoice_pdf": inv.invoice_pdf,
|
||||||
|
}
|
||||||
|
for inv in invoices.auto_paging_iter()
|
||||||
|
]
|
||||||
|
except Exception:
|
||||||
|
return []
|
||||||
|
|
||||||
# ── Private DB helpers ───────────────────────────────────────────────
|
# ── Private DB helpers ───────────────────────────────────────────────
|
||||||
|
|
||||||
async def _upsert_subscription(
|
async def _upsert_subscription(
|
||||||
|
|||||||
@@ -21,42 +21,58 @@ FEATURES: dict[str, dict[str, Any]] = {
|
|||||||
"free": {
|
"free": {
|
||||||
"agents": 3,
|
"agents": 3,
|
||||||
"batch_active": 2,
|
"batch_active": 2,
|
||||||
"cloud_storage_gb": 0,
|
"batch_runs_per_day": 5,
|
||||||
"backup_gb": 0,
|
|
||||||
"providers": 1,
|
"providers": 1,
|
||||||
"batch_builder": False,
|
"batch_builder": False,
|
||||||
"plugin_marketplace": False,
|
|
||||||
"sso": False,
|
"sso": False,
|
||||||
|
"real_embeddings": False, # keyword fallback only
|
||||||
|
"realtime_extraction": False, # batch queue (Phase 2)
|
||||||
|
"relational_memory": False, # relational tier (Phase 3) — Pro+
|
||||||
|
"proactive_mining": False, # Power+ only (Phase 5)
|
||||||
|
"folder_max_files": 200,
|
||||||
|
"folder_monthly_tokens": 100_000,
|
||||||
},
|
},
|
||||||
"pro": {
|
"pro": {
|
||||||
"agents": -1, # unlimited
|
"agents": -1, # unlimited
|
||||||
"batch_active": 10,
|
"batch_active": 10,
|
||||||
"cloud_storage_gb": 5,
|
"batch_runs_per_day": 50,
|
||||||
"backup_gb": 5,
|
|
||||||
"providers": -1,
|
"providers": -1,
|
||||||
"batch_builder": False,
|
"batch_builder": False,
|
||||||
"plugin_marketplace": False,
|
|
||||||
"sso": False,
|
"sso": False,
|
||||||
|
"real_embeddings": True, # pgvector cosine search
|
||||||
|
"realtime_extraction": True, # fire-and-forget asyncio.create_task
|
||||||
|
"relational_memory": True, # person/project predicates
|
||||||
|
"proactive_mining": False, # Power+ only (Phase 5)
|
||||||
|
"folder_max_files": 5000,
|
||||||
|
"folder_monthly_tokens": 2_000_000,
|
||||||
},
|
},
|
||||||
"power": {
|
"power": {
|
||||||
"agents": -1,
|
"agents": -1,
|
||||||
"batch_active": -1, # unlimited
|
"batch_active": -1, # unlimited
|
||||||
"cloud_storage_gb": 25,
|
"batch_runs_per_day": -1, # unlimited
|
||||||
"backup_gb": 25,
|
|
||||||
"providers": -1,
|
"providers": -1,
|
||||||
"batch_builder": True,
|
"batch_builder": True,
|
||||||
"plugin_marketplace": True,
|
|
||||||
"sso": False,
|
"sso": False,
|
||||||
|
"real_embeddings": True,
|
||||||
|
"realtime_extraction": True,
|
||||||
|
"relational_memory": True, # all predicates incl. custom
|
||||||
|
"proactive_mining": True, # scheduled pattern mining (Phase 5)
|
||||||
|
"folder_max_files": -1, # unlimited
|
||||||
|
"folder_monthly_tokens": -1, # unlimited
|
||||||
},
|
},
|
||||||
"team": {
|
"team": {
|
||||||
"agents": -1,
|
"agents": -1,
|
||||||
"batch_active": -1,
|
"batch_active": -1,
|
||||||
"cloud_storage_gb": -1, # unlimited
|
"batch_runs_per_day": -1, # unlimited
|
||||||
"backup_gb": -1, # unlimited
|
|
||||||
"providers": -1,
|
"providers": -1,
|
||||||
"batch_builder": True,
|
"batch_builder": True,
|
||||||
"plugin_marketplace": True,
|
|
||||||
"sso": True,
|
"sso": True,
|
||||||
|
"real_embeddings": True,
|
||||||
|
"realtime_extraction": True,
|
||||||
|
"relational_memory": True, # all predicates incl. custom
|
||||||
|
"proactive_mining": True, # scheduled pattern mining (Phase 5)
|
||||||
|
"folder_max_files": -1, # unlimited
|
||||||
|
"folder_monthly_tokens": -1, # unlimited
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -77,16 +93,18 @@ class TierManager:
|
|||||||
async def get_tier(self, user_id: str, db: AsyncSession) -> BillingTier:
|
async def get_tier(self, user_id: str, db: AsyncSession) -> BillingTier:
|
||||||
"""Return the current billing tier for ``user_id`` from the DB.
|
"""Return the current billing tier for ``user_id`` from the DB.
|
||||||
|
|
||||||
Falls back to ``'free'`` when no subscription row exists.
|
Falls back to ``'power'`` in dev (unlimited) or ``'free'`` in prod
|
||||||
|
when no subscription row exists.
|
||||||
"""
|
"""
|
||||||
from app.models import Subscription # noqa: PLC0415
|
from app.models import Subscription # noqa: PLC0415
|
||||||
|
from app.config.settings import settings # noqa: PLC0415
|
||||||
|
|
||||||
result = await db.execute(
|
result = await db.execute(
|
||||||
select(Subscription.tier).where(Subscription.user_id == user_id)
|
select(Subscription.tier).where(Subscription.user_id == user_id)
|
||||||
)
|
)
|
||||||
tier: str | None = result.scalar_one_or_none()
|
tier: str | None = result.scalar_one_or_none()
|
||||||
if tier is None or tier not in FEATURES:
|
if tier is None or tier not in FEATURES:
|
||||||
return "free"
|
return "power" if settings.ENV == "dev" else "free"
|
||||||
return tier # type: ignore[return-value]
|
return tier # type: ignore[return-value]
|
||||||
|
|
||||||
# ── Feature access ───────────────────────────────────────────────────
|
# ── Feature access ───────────────────────────────────────────────────
|
||||||
@@ -113,77 +131,19 @@ class TierManager:
|
|||||||
)
|
)
|
||||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=detail)
|
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=detail)
|
||||||
|
|
||||||
|
def get_feature_value(self, tier: BillingTier, feature: str) -> int:
|
||||||
|
"""Return integer feature value for tier. -1 means unlimited."""
|
||||||
|
value = FEATURES.get(tier, FEATURES["free"]).get(feature)
|
||||||
|
if not isinstance(value, int):
|
||||||
|
return 0
|
||||||
|
return value
|
||||||
|
|
||||||
# ── Rate limiting ────────────────────────────────────────────────────
|
# ── Rate limiting ────────────────────────────────────────────────────
|
||||||
|
|
||||||
def get_rate_limit(self, tier: BillingTier) -> int:
|
def get_rate_limit(self, tier: BillingTier) -> int:
|
||||||
"""Return the requests-per-minute limit for ``tier``."""
|
"""Return the requests-per-minute limit for ``tier``."""
|
||||||
return RATE_LIMITS.get(tier, RATE_LIMITS["free"])
|
return RATE_LIMITS.get(tier, RATE_LIMITS["free"])
|
||||||
|
|
||||||
# ── Storage quota ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def enforce_quota(
|
|
||||||
self,
|
|
||||||
tier: BillingTier,
|
|
||||||
current_bytes: int = 0,
|
|
||||||
additional_bytes: int = 0,
|
|
||||||
) -> None:
|
|
||||||
"""Raise ``HTTP 402`` if the user would exceed their cloud storage quota.
|
|
||||||
|
|
||||||
``tier`` is the caller's current tier (from ``current_user.tier``).
|
|
||||||
``current_bytes`` is the total bytes already stored (queried by caller).
|
|
||||||
"""
|
|
||||||
limit_gb: int = FEATURES[tier]["cloud_storage_gb"]
|
|
||||||
if limit_gb == 0:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
|
||||||
detail=f"Cloud storage is not available on the '{tier}' tier",
|
|
||||||
)
|
|
||||||
if limit_gb == -1:
|
|
||||||
return # unlimited
|
|
||||||
limit_bytes = limit_gb * 1024 ** 3
|
|
||||||
if current_bytes + additional_bytes > limit_bytes:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
|
||||||
detail=f"Storage quota exceeded for tier '{tier}'",
|
|
||||||
)
|
|
||||||
|
|
||||||
def enforce_backup_quota(
|
|
||||||
self,
|
|
||||||
tier: BillingTier,
|
|
||||||
current_bytes: int = 0,
|
|
||||||
additional_bytes: int = 0,
|
|
||||||
) -> None:
|
|
||||||
"""Raise ``HTTP 402`` if the user would exceed their backup quota."""
|
|
||||||
limit_gb: int = FEATURES[tier]["backup_gb"]
|
|
||||||
if limit_gb == 0:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
|
||||||
detail=f"Backup is not available on the '{tier}' tier",
|
|
||||||
)
|
|
||||||
if limit_gb == -1:
|
|
||||||
return # unlimited
|
|
||||||
limit_bytes = limit_gb * 1024 ** 3
|
|
||||||
if current_bytes + additional_bytes > limit_bytes:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
|
||||||
detail=f"Backup quota exceeded for tier '{tier}'",
|
|
||||||
)
|
|
||||||
|
|
||||||
def check_quota(
|
|
||||||
self,
|
|
||||||
tier: BillingTier,
|
|
||||||
current_bytes: int = 0,
|
|
||||||
additional_bytes: int = 0,
|
|
||||||
) -> bool:
|
|
||||||
"""Return ``True`` if the user can store ``additional_bytes`` more data."""
|
|
||||||
limit_gb: int = FEATURES[tier]["cloud_storage_gb"]
|
|
||||||
if limit_gb == 0:
|
|
||||||
return False
|
|
||||||
if limit_gb == -1:
|
|
||||||
return True
|
|
||||||
limit_bytes = limit_gb * 1024 ** 3
|
|
||||||
return current_bytes + additional_bytes <= limit_bytes
|
|
||||||
|
|
||||||
|
|
||||||
# Module-level singleton shared across the app.
|
# Module-level singleton shared across the app.
|
||||||
tier_manager = TierManager()
|
tier_manager = TierManager()
|
||||||
|
|||||||
@@ -1,9 +1,9 @@
|
|||||||
from typing import Literal
|
from typing import Literal
|
||||||
from pydantic_settings import BaseSettings
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||||
|
|
||||||
|
|
||||||
class Settings(BaseSettings):
|
class Settings(BaseSettings):
|
||||||
DATABASE_URL: str = "postgresql+asyncpg://postgres:postgres@localhost:5432/adiuva"
|
DATABASE_URL: str = "postgresql+asyncpg://postgres:postgres@localhost:5432/adiuvai"
|
||||||
JWT_SECRET: str = "change-me-in-production"
|
JWT_SECRET: str = "change-me-in-production"
|
||||||
JWT_ALGORITHM: str = "HS256"
|
JWT_ALGORITHM: str = "HS256"
|
||||||
JWT_ACCESS_TOKEN_EXPIRE_MINUTES: int = 30
|
JWT_ACCESS_TOKEN_EXPIRE_MINUTES: int = 30
|
||||||
@@ -12,31 +12,75 @@ class Settings(BaseSettings):
|
|||||||
STRIPE_SECRET_KEY: str = ""
|
STRIPE_SECRET_KEY: str = ""
|
||||||
STRIPE_WEBHOOK_SECRET: str = ""
|
STRIPE_WEBHOOK_SECRET: str = ""
|
||||||
|
|
||||||
S3_BUCKET: str = ""
|
|
||||||
S3_REGION: str = "us-east-1"
|
|
||||||
S3_ENDPOINT_URL: str = ""
|
|
||||||
AWS_ACCESS_KEY_ID: str = ""
|
|
||||||
AWS_SECRET_ACCESS_KEY: str = ""
|
|
||||||
|
|
||||||
PINECONE_API_KEY: str = ""
|
|
||||||
PINECONE_INDEX: str = "adiuva"
|
|
||||||
QDRANT_URL: str = ""
|
|
||||||
QDRANT_API_KEY: str = ""
|
|
||||||
|
|
||||||
OPENAI_API_KEY: str = ""
|
OPENAI_API_KEY: str = ""
|
||||||
ANTHROPIC_API_KEY: str = ""
|
ANTHROPIC_API_KEY: str = ""
|
||||||
GOOGLE_API_KEY: str = ""
|
GOOGLE_API_KEY: str = ""
|
||||||
|
CEREBRAS_API_KEY: str = ""
|
||||||
|
GROQ_API_KEY: str = ""
|
||||||
|
DEEPSEEK_API_KEY: str = ""
|
||||||
|
|
||||||
LLM_MODEL: str = "gpt-4o"
|
LLM_MODEL: str = "gpt-4o"
|
||||||
LLM_ROUTER_MODEL: str = "gpt-4o-mini"
|
LLM_EMBED_MODEL: str = "text-embedding-3-small"
|
||||||
|
|
||||||
CORS_ORIGINS: list[str] = ["app://.", "http://localhost:3000", "http://localhost:5173"]
|
# Per-agent model overrides. Leave empty to fall back to LLM_MODEL.
|
||||||
|
LLM_MODEL_CLASSIFIER: str = "" # _infer_floating_domain (intent routing)
|
||||||
|
LLM_MODEL_HOME_AGENT: str = "" # home-agent (run_single_agent / stream)
|
||||||
|
LLM_MODEL_FLOATING_AGENT: str = "" # floating-agent (contextual chat)
|
||||||
|
LLM_MODEL_UNIFIED_PROCESSOR: str = "" # unified-processor (agent_runner)
|
||||||
|
LLM_MODEL_CLOUD_PROCESSOR: str = "" # cloud-processor (agent_runner)
|
||||||
|
LLM_MODEL_BRIEF_AGENT: str = "" # brief-agent (home + project text briefs)
|
||||||
|
LLM_MODEL_TASK_BRIEF_AGENT: str = "" # task-brief-agent (per-task deep research)
|
||||||
|
LLM_MODEL_SETUP_AGENT: str = "" # agent-setup journey
|
||||||
|
LLM_MODEL_MEMORY_EXTRACTOR: str = "" # memory-extractor (Phase 2 extract/decide)
|
||||||
|
LLM_MODEL_MEMORY_MINER: str = "" # memory-miner (Phase 5 proactive mining)
|
||||||
|
LLM_MODEL_MEMORY_AUDITOR: str = "" # memory-auditor (Phase 7 weekly audit)
|
||||||
|
|
||||||
|
# GitHub Copilot OAuth token storage directory.
|
||||||
|
# Leave empty to use the LiteLLM default (~/.config/litellm/github_copilot).
|
||||||
|
# In Docker, set this to a path backed by a named volume so tokens survive restarts.
|
||||||
|
GITHUB_COPILOT_TOKEN_DIR: str = ""
|
||||||
|
|
||||||
|
# OAuth client credentials — used for Gmail and Microsoft (Outlook/Teams) flows.
|
||||||
|
GMAIL_CLIENT_ID: str = ""
|
||||||
|
GMAIL_CLIENT_SECRET: str = ""
|
||||||
|
MS_CLIENT_ID: str = ""
|
||||||
|
MS_CLIENT_SECRET: str = ""
|
||||||
|
# MS_TENANT_ID: set to 'common' to allow multi-tenant (personal + work accounts).
|
||||||
|
MS_TENANT_ID: str = "common"
|
||||||
|
|
||||||
|
# Google Login OAuth credentials — scope: openid email profile.
|
||||||
|
# Separate from GMAIL_CLIENT_ID/SECRET (which uses gmail.readonly scope).
|
||||||
|
GOOGLE_AUTH_CLIENT_ID: str = ""
|
||||||
|
GOOGLE_AUTH_CLIENT_SECRET: str = ""
|
||||||
|
# The redirect URI registered in Google Cloud Console.
|
||||||
|
# Google redirects here after consent; this backend route then bounces to
|
||||||
|
# the adiuvai:// deep link so the Electron app receives the code.
|
||||||
|
# Dev: http://localhost:8000/api/v1/auth/oauth/google/web-callback
|
||||||
|
# Prod: https://api.adiuvai.com/api/v1/auth/oauth/google/web-callback
|
||||||
|
OAUTH_REDIRECT_URI: str = "http://localhost:8000/api/v1/auth/oauth/google/web-callback"
|
||||||
|
|
||||||
|
# Fernet key (URL-safe base64, 32-byte key) for at-rest encryption of OAuth
|
||||||
|
# tokens stored in cloud_agent_configs.oauth_token_encrypted.
|
||||||
|
# Generate with: from cryptography.fernet import Fernet; Fernet.generate_key()
|
||||||
|
OAUTH_ENCRYPTION_KEY: str = ""
|
||||||
|
|
||||||
|
CORS_ORIGINS: list[str] = [
|
||||||
|
"app://.",
|
||||||
|
"http://localhost:3000",
|
||||||
|
"http://localhost:5173",
|
||||||
|
"http://localhost:4173", # Vite preview (web SPA)
|
||||||
|
"https://app.adiuvai.com", # Production web portal
|
||||||
|
]
|
||||||
|
|
||||||
|
LANGFUSE_SECRET_KEY: str = ""
|
||||||
|
LANGFUSE_PUBLIC_KEY: str = ""
|
||||||
|
LANGFUSE_BASE_URL: str = "https://cloud.langfuse.com"
|
||||||
|
|
||||||
|
SCHEDULER_ENABLED: bool = True
|
||||||
|
|
||||||
ENV: Literal["dev", "prod"] = "dev"
|
ENV: Literal["dev", "prod"] = "dev"
|
||||||
|
|
||||||
class Config:
|
model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8", extra="ignore")
|
||||||
env_file = ".env"
|
|
||||||
env_file_encoding = "utf-8"
|
|
||||||
|
|
||||||
|
|
||||||
settings = Settings()
|
settings = Settings()
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
"""Agent Registry — base classes and singleton registry for chat agents."""
|
"""Minimal agent base types retained for compatibility with batch runners."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
@@ -7,7 +7,7 @@ from typing import Any
|
|||||||
|
|
||||||
|
|
||||||
class BaseAgent(ABC):
|
class BaseAgent(ABC):
|
||||||
"""Common base for all agents."""
|
"""Common base for non-chat agents still using the old base contract."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -27,111 +27,4 @@ class BaseAgent(ABC):
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def skills(self) -> list[str]:
|
def skills(self) -> list[str]:
|
||||||
"""Override in subclasses to advertise capabilities."""
|
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
||||||
class ChatAgent(BaseAgent):
|
|
||||||
"""Base class for LLM-powered chat agents."""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
|
||||||
"""Process a user query and return a text response."""
|
|
||||||
...
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
def get_tools(self) -> list[Any]:
|
|
||||||
"""Return LangChain tool definitions available to this agent."""
|
|
||||||
...
|
|
||||||
|
|
||||||
async def _tool_loop(
|
|
||||||
self,
|
|
||||||
llm: Any,
|
|
||||||
messages: list[Any],
|
|
||||||
tools: list[Any],
|
|
||||||
max_iter: int = 5,
|
|
||||||
) -> str:
|
|
||||||
"""Shared tool-calling loop.
|
|
||||||
|
|
||||||
Binds *tools* to *llm*, invokes iteratively until the model stops
|
|
||||||
requesting tool calls or *max_iter* is reached, and returns the
|
|
||||||
final text response.
|
|
||||||
"""
|
|
||||||
from langchain_core.messages import AIMessage, ToolMessage
|
|
||||||
|
|
||||||
llm_with_tools = llm.bind_tools(tools) if tools else llm
|
|
||||||
|
|
||||||
for _ in range(max_iter):
|
|
||||||
response: AIMessage = await llm_with_tools.ainvoke(messages)
|
|
||||||
messages.append(response)
|
|
||||||
|
|
||||||
if not response.tool_calls:
|
|
||||||
return str(response.content)
|
|
||||||
|
|
||||||
# Execute each requested tool call
|
|
||||||
tool_map = {t.name: t for t in tools}
|
|
||||||
for call in response.tool_calls:
|
|
||||||
tool_fn = tool_map.get(call["name"])
|
|
||||||
if tool_fn is None:
|
|
||||||
result = f"Unknown tool: {call['name']}"
|
|
||||||
else:
|
|
||||||
result = await tool_fn.ainvoke(call["args"])
|
|
||||||
messages.append(
|
|
||||||
ToolMessage(content=str(result), tool_call_id=call["id"])
|
|
||||||
)
|
|
||||||
|
|
||||||
# Exhausted iterations — ask model for a final answer without tools
|
|
||||||
response = await llm.ainvoke(messages)
|
|
||||||
return str(response.content)
|
|
||||||
|
|
||||||
|
|
||||||
class AgentRegistry:
|
|
||||||
"""Singleton registry for ChatAgent subclasses."""
|
|
||||||
|
|
||||||
_instance: AgentRegistry | None = None
|
|
||||||
|
|
||||||
def __init__(self) -> None:
|
|
||||||
self._agents: dict[str, type[ChatAgent]] = {}
|
|
||||||
|
|
||||||
def __new__(cls) -> AgentRegistry:
|
|
||||||
if cls._instance is None:
|
|
||||||
cls._instance = super().__new__(cls)
|
|
||||||
cls._instance._agents = {}
|
|
||||||
return cls._instance
|
|
||||||
|
|
||||||
# ── public API ───────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def register(self, agent_class: type[ChatAgent]) -> type[ChatAgent]:
|
|
||||||
"""Class decorator — registers an agent by its name."""
|
|
||||||
instance = agent_class()
|
|
||||||
name = instance.get_name()
|
|
||||||
self._agents[name] = agent_class
|
|
||||||
return agent_class
|
|
||||||
|
|
||||||
def get(self, name: str) -> ChatAgent:
|
|
||||||
"""Return a fresh instance of the named agent."""
|
|
||||||
cls = self._agents.get(name)
|
|
||||||
if cls is None:
|
|
||||||
raise KeyError(f"Agent not found: {name}")
|
|
||||||
return cls()
|
|
||||||
|
|
||||||
def list_agents(self) -> list[dict[str, str]]:
|
|
||||||
"""Return ``[{name, description}]`` for the orchestrator prompt."""
|
|
||||||
result: list[dict[str, str]] = []
|
|
||||||
for cls in self._agents.values():
|
|
||||||
inst = cls()
|
|
||||||
result.append(
|
|
||||||
{"name": inst.get_name(), "description": inst.get_description()}
|
|
||||||
)
|
|
||||||
return result
|
|
||||||
|
|
||||||
async def call_agent(
|
|
||||||
self, name: str, query: str, context: dict[str, Any]
|
|
||||||
) -> str:
|
|
||||||
"""Instantiate the named agent and call its ``handle`` method."""
|
|
||||||
agent = self.get(name)
|
|
||||||
return await agent.handle(query, context)
|
|
||||||
|
|
||||||
|
|
||||||
# Module-level singleton
|
|
||||||
registry = AgentRegistry()
|
|
||||||
|
|||||||
1051
app/core/agent_runner.py
Normal file
1051
app/core/agent_runner.py
Normal file
File diff suppressed because it is too large
Load Diff
59
app/core/agent_session_buffer.py
Normal file
59
app/core/agent_session_buffer.py
Normal file
@@ -0,0 +1,59 @@
|
|||||||
|
"""In-process TTL buffer for per-session LangChain message history.
|
||||||
|
|
||||||
|
Stores the full message list (including AIMessage with tool_calls and ToolMessage)
|
||||||
|
keyed by (user_id, session_id), so agents can reconstruct tool-call context across
|
||||||
|
conversation turns without it being lossy through the wire.
|
||||||
|
|
||||||
|
Single-process only. For multi-worker deployments, replace the _SessionBuffer
|
||||||
|
implementation with one backed by Redis (serialize LangChain messages to dicts via
|
||||||
|
message_to_dict / messages_from_dict from langchain_core.messages).
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from threading import Lock
|
||||||
|
|
||||||
|
from langchain_core.messages import BaseMessage
|
||||||
|
|
||||||
|
SESSION_TTL_SECONDS = 1800 # 30-minute idle expiry
|
||||||
|
MAX_MESSAGES_PER_SESSION = 80 # cap to avoid unbounded memory growth
|
||||||
|
|
||||||
|
|
||||||
|
class _SessionBuffer:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._store: dict[tuple[str, str], tuple[float, list[BaseMessage]]] = {}
|
||||||
|
self._lock = Lock()
|
||||||
|
|
||||||
|
def _evict_stale(self) -> None:
|
||||||
|
now = time.monotonic()
|
||||||
|
stale = [k for k, (ts, _) in self._store.items() if now - ts > SESSION_TTL_SECONDS]
|
||||||
|
for k in stale:
|
||||||
|
del self._store[k]
|
||||||
|
|
||||||
|
def get(self, user_id: str, session_id: str) -> list[BaseMessage] | None:
|
||||||
|
key = (user_id, session_id)
|
||||||
|
with self._lock:
|
||||||
|
entry = self._store.get(key)
|
||||||
|
if entry is None:
|
||||||
|
return None
|
||||||
|
ts, msgs = entry
|
||||||
|
if time.monotonic() - ts > SESSION_TTL_SECONDS:
|
||||||
|
del self._store[key]
|
||||||
|
return None
|
||||||
|
self._store[key] = (time.monotonic(), msgs)
|
||||||
|
return list(msgs)
|
||||||
|
|
||||||
|
def set(self, user_id: str, session_id: str, messages: list[BaseMessage]) -> None:
|
||||||
|
key = (user_id, session_id)
|
||||||
|
capped = messages[-MAX_MESSAGES_PER_SESSION:]
|
||||||
|
with self._lock:
|
||||||
|
self._evict_stale()
|
||||||
|
self._store[key] = (time.monotonic(), capped)
|
||||||
|
|
||||||
|
def clear(self, user_id: str, session_id: str) -> None:
|
||||||
|
with self._lock:
|
||||||
|
self._store.pop((user_id, session_id), None)
|
||||||
|
|
||||||
|
|
||||||
|
# Module-level singleton — same pattern as _pending_states in api/app/api/routes/auth.py
|
||||||
|
session_buffer = _SessionBuffer()
|
||||||
228
app/core/brief_agent.py
Normal file
228
app/core/brief_agent.py
Normal file
@@ -0,0 +1,228 @@
|
|||||||
|
"""Brief agent — produces plain-text home and project status briefs.
|
||||||
|
|
||||||
|
Read-only tool subset only. Never calls _normalize_tagged_list_lines —
|
||||||
|
the brief prompt forbids XML tags, so skipping post-processing is intentional.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
from datetime import date
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from app.agents.note_agent import NOTE_READ_TOOLS
|
||||||
|
from app.agents.project_agent import PROJECT_READ_TOOLS
|
||||||
|
from app.agents.task_agent import TASK_READ_TOOLS
|
||||||
|
from app.agents.timeline_agent import TIMELINE_READ_TOOLS
|
||||||
|
from app.core.deep_agent import (
|
||||||
|
_language_instruction,
|
||||||
|
_proactive_hints_injection,
|
||||||
|
_read_only_memory_tools,
|
||||||
|
_relational_memory_injection,
|
||||||
|
_run_single_agent_stream,
|
||||||
|
_trace_id_from_context,
|
||||||
|
build_brief_multi_project_manifest,
|
||||||
|
)
|
||||||
|
from app.core.langfuse_client import compile_prompt, get_prompt_or_fallback
|
||||||
|
|
||||||
|
_LANGUAGE_NAMES: dict[str, str] = {
|
||||||
|
"en": "English", "it": "Italian", "es": "Spanish",
|
||||||
|
"fr": "French", "de": "German",
|
||||||
|
"english": "English", "italian": "Italian", "italiano": "Italian",
|
||||||
|
"spanish": "Spanish", "español": "Spanish",
|
||||||
|
"french": "French", "français": "French",
|
||||||
|
"german": "German", "deutsch": "German",
|
||||||
|
}
|
||||||
|
|
||||||
|
_HOME_BRIEF_FALLBACK = """\
|
||||||
|
You are the user's personal assistant producing a short daily brief.
|
||||||
|
|
||||||
|
ROLE
|
||||||
|
Act like a calm, attentive secretary writing a stand-up note for your boss.
|
||||||
|
Warm and human, never breezy. Never cheerful filler, never emojis, never
|
||||||
|
"here is your brief" meta-text. The user is opening the app mid-workday and
|
||||||
|
is probably stressed — your job is to lower cognitive load, not add noise.
|
||||||
|
|
||||||
|
TOOLS — always call before writing
|
||||||
|
Pull fresh data every run. Do not invent counts or titles. Use at minimum:
|
||||||
|
- list_tasks_due_today — tasks the user owes today
|
||||||
|
- list_timelines_today — events starting or ending today
|
||||||
|
- list_all_projects — projects currently in progress or at risk
|
||||||
|
- memory_list_blocks / memory_get — personal context about people, clients,
|
||||||
|
payment habits, working preferences
|
||||||
|
If a tool returns nothing, simply omit that topic. Never report zeros.
|
||||||
|
|
||||||
|
WHAT TO INCLUDE
|
||||||
|
1. Tasks due today (title + priority; group the 1-2 most important).
|
||||||
|
2. Timeline events starting or ending today (and anything that starts/ends
|
||||||
|
tomorrow if the user has a very light day).
|
||||||
|
3. Active projects that need a nudge — stalled, blocked, or awaiting input.
|
||||||
|
4. Memory-aware colour where it sharpens the brief. Examples:
|
||||||
|
- "Client Rossi tends to pay late — the Acme invoice is 6 days out."
|
||||||
|
- "You usually dislike meetings before 10:00 — the call at 09:30 is unusual."
|
||||||
|
Only add a memory line when it changes what the user does. Do not pad.
|
||||||
|
|
||||||
|
WHAT TO OMIT
|
||||||
|
- Zero-counts ("no overdue items", "0 meetings today").
|
||||||
|
- Statistics ("2 active projects, 3 completed tasks").
|
||||||
|
- Headers, titles, greetings, sign-offs, dates, emojis, slang.
|
||||||
|
- Meta-phrases ("here is", "let me know if", "hope this helps").
|
||||||
|
- XML/HTML tags of any kind. Plain prose only.
|
||||||
|
|
||||||
|
LIGHT-DAY CLAUSE
|
||||||
|
If tasks + events + active-project-nudges together produce fewer than two
|
||||||
|
sentences of content, also list 1-2 projects in status on_hold or waiting
|
||||||
|
and ask a single, specific question about them — e.g. "Is the Bianchi
|
||||||
|
redesign still paused, or ready to pick back up?" One question max, grounded
|
||||||
|
in a real project name.
|
||||||
|
|
||||||
|
VOICE
|
||||||
|
- Calm. Concise. Human. Short sentences.
|
||||||
|
- Use **bold** sparingly for task titles, project names, and people's names.
|
||||||
|
- No bullet lists. Flow as 2-4 sentences of prose.
|
||||||
|
|
||||||
|
LENGTH
|
||||||
|
2-4 sentences total. Hard cap 4. If the day is truly empty, one sentence.
|
||||||
|
|
||||||
|
Respond in the user's language ({language}). Today is {today}.\
|
||||||
|
"""
|
||||||
|
|
||||||
|
_PROJECT_BRIEF_FALLBACK = """\
|
||||||
|
You are the project assistant producing a short status brief for ONE project.
|
||||||
|
|
||||||
|
ROLE
|
||||||
|
A senior project manager summarising state-of-play for the owner. Factual,
|
||||||
|
sharp, forward-looking. Never reassuring filler, never emojis.
|
||||||
|
|
||||||
|
SCOPE
|
||||||
|
Work only with project_id = {project_id}. Do not mention or pull data from
|
||||||
|
other projects. Use tools to fetch fresh data:
|
||||||
|
- get_project — current status, dates, description
|
||||||
|
- list_tasks(project_id) — open work, split by status
|
||||||
|
- list_timelines(project_id) — milestones hit, upcoming, overdue
|
||||||
|
- list_notes(project_id) — any recent decisions or blockers
|
||||||
|
- memory_get — relevant context about the client, collaborators, constraints
|
||||||
|
|
||||||
|
STRUCTURE — follow exactly, one short paragraph per section, no headers
|
||||||
|
1. **State.** One sentence: current phase, health (on track / at risk / blocked),
|
||||||
|
and why. Cite the concrete signal (overdue milestone, stalled tasks, recent
|
||||||
|
blocker note).
|
||||||
|
2. **What's moving.** What was completed or progressed recently. Name specific
|
||||||
|
tasks or milestones.
|
||||||
|
3. **Next steps.** The 1-3 most important things the user should do next, in
|
||||||
|
priority order. Be concrete — task name, who owns it, when due if known.
|
||||||
|
If waiting on someone else, name them and what the ask is.
|
||||||
|
4. **Risks / memory-flagged items.** One line max. Only include when there is
|
||||||
|
a real risk or a relevant memory (e.g. late-paying client, tight deadline,
|
||||||
|
scope change). Omit the section entirely if nothing to say.
|
||||||
|
|
||||||
|
WHAT TO OMIT
|
||||||
|
- Zero-counts ("no overdue tasks").
|
||||||
|
- Generic advice ("keep up the good work").
|
||||||
|
- Greetings, headers, bullet lists, emojis, sign-offs, meta-phrases.
|
||||||
|
- XML/HTML tags or bracketed id lists. Plain prose only.
|
||||||
|
|
||||||
|
VOICE
|
||||||
|
- Direct. Factual. No fluff.
|
||||||
|
- Use **bold** sparingly for task titles, milestone names, and the owner's name.
|
||||||
|
- Short sentences. Prefer verbs over nouns ("Client review is blocking release"
|
||||||
|
not "There is a blocker which is the client review").
|
||||||
|
|
||||||
|
LENGTH
|
||||||
|
4-8 sentences total across the 3-4 sections. Hard cap 8.
|
||||||
|
|
||||||
|
Respond in the user's language ({language}). Today is {today}.\
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_language(context: dict[str, Any]) -> str:
|
||||||
|
core = context.get("core_memory") or {}
|
||||||
|
raw = (core.get("language") or "en").strip().lower()
|
||||||
|
return _LANGUAGE_NAMES.get(raw, raw.title()) or "English"
|
||||||
|
|
||||||
|
|
||||||
|
def _build_read_tools(user_id: str, trace_id: str | None) -> list[Any]:
|
||||||
|
return [
|
||||||
|
*TASK_READ_TOOLS,
|
||||||
|
*PROJECT_READ_TOOLS,
|
||||||
|
*TIMELINE_READ_TOOLS,
|
||||||
|
*NOTE_READ_TOOLS,
|
||||||
|
*_read_only_memory_tools(user_id, trace_id),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
async def run_home_brief(
|
||||||
|
user_id: str,
|
||||||
|
context: dict[str, Any],
|
||||||
|
) -> AsyncGenerator[tuple[str, Any], None]:
|
||||||
|
"""Stream a plain-text daily home brief.
|
||||||
|
|
||||||
|
Yields (event_type, data) tuples identical to _run_single_agent_stream.
|
||||||
|
Do NOT post-process output through _normalize_tagged_list_lines.
|
||||||
|
"""
|
||||||
|
from app.agents.folder_agent import FOLDER_TOOLS
|
||||||
|
|
||||||
|
trace_id = _trace_id_from_context(context)
|
||||||
|
today = date.today().isoformat()
|
||||||
|
language = _resolve_language(context)
|
||||||
|
|
||||||
|
raw_template, langfuse_prompt = get_prompt_or_fallback("home_brief", _HOME_BRIEF_FALLBACK)
|
||||||
|
system_prompt = compile_prompt(raw_template, langfuse_prompt, language=language, today=today)
|
||||||
|
system_prompt += _relational_memory_injection(context)
|
||||||
|
system_prompt += _proactive_hints_injection(context)
|
||||||
|
system_prompt += _language_instruction(context)
|
||||||
|
if today not in system_prompt:
|
||||||
|
system_prompt += f"\nToday is {today}."
|
||||||
|
|
||||||
|
brief_manifest = await build_brief_multi_project_manifest()
|
||||||
|
system_prompt = system_prompt + ("\n\n" + brief_manifest if brief_manifest else "")
|
||||||
|
|
||||||
|
tools = [*_build_read_tools(user_id, trace_id), *FOLDER_TOOLS]
|
||||||
|
async for event in _run_single_agent_stream(
|
||||||
|
user_id=user_id,
|
||||||
|
system_prompt=system_prompt,
|
||||||
|
message="Generate the daily brief.",
|
||||||
|
context=context,
|
||||||
|
langfuse_prompt=langfuse_prompt,
|
||||||
|
agent_name="brief-agent",
|
||||||
|
tools=tools,
|
||||||
|
):
|
||||||
|
yield event
|
||||||
|
|
||||||
|
|
||||||
|
async def run_project_brief(
|
||||||
|
user_id: str,
|
||||||
|
project_id: str,
|
||||||
|
context: dict[str, Any],
|
||||||
|
) -> AsyncGenerator[tuple[str, Any], None]:
|
||||||
|
"""Stream a plain-text project status brief for project_id.
|
||||||
|
|
||||||
|
Yields (event_type, data) tuples identical to _run_single_agent_stream.
|
||||||
|
Do NOT post-process output through _normalize_tagged_list_lines.
|
||||||
|
"""
|
||||||
|
trace_id = _trace_id_from_context(context)
|
||||||
|
today = date.today().isoformat()
|
||||||
|
language = _resolve_language(context)
|
||||||
|
|
||||||
|
raw_template, langfuse_prompt = get_prompt_or_fallback("project_brief", _PROJECT_BRIEF_FALLBACK)
|
||||||
|
system_prompt = compile_prompt(
|
||||||
|
raw_template, langfuse_prompt,
|
||||||
|
language=language, today=today, project_id=project_id,
|
||||||
|
)
|
||||||
|
system_prompt += _relational_memory_injection(context)
|
||||||
|
system_prompt += _proactive_hints_injection(context)
|
||||||
|
system_prompt += _language_instruction(context)
|
||||||
|
if today not in system_prompt:
|
||||||
|
system_prompt += f"\nToday is {today}."
|
||||||
|
|
||||||
|
tools = _build_read_tools(user_id, trace_id)
|
||||||
|
async for event in _run_single_agent_stream(
|
||||||
|
user_id=user_id,
|
||||||
|
system_prompt=system_prompt,
|
||||||
|
message=f"Generate the project status brief for project {project_id}.",
|
||||||
|
context=context,
|
||||||
|
langfuse_prompt=langfuse_prompt,
|
||||||
|
agent_name="brief-agent",
|
||||||
|
tools=tools,
|
||||||
|
):
|
||||||
|
yield event
|
||||||
1581
app/core/deep_agent.py
Normal file
1581
app/core/deep_agent.py
Normal file
File diff suppressed because it is too large
Load Diff
151
app/core/device_manager.py
Normal file
151
app/core/device_manager.py
Normal file
@@ -0,0 +1,151 @@
|
|||||||
|
"""Device connection manager.
|
||||||
|
|
||||||
|
Maintains in-memory state for all active Electron → backend WebSocket
|
||||||
|
connections. One connection per user (latest replaces previous).
|
||||||
|
|
||||||
|
The manager handles the **tool-call round-trip** pattern:
|
||||||
|
- Backend sends ``tool_call`` frame → Electron executes the action →
|
||||||
|
returns ``tool_result`` frame.
|
||||||
|
- ``create_pending_call`` registers a Future keyed by ``call_id``.
|
||||||
|
- ``resolve_pending_call`` fulfils the Future; callers awaiting it
|
||||||
|
receive the result dict from Electron.
|
||||||
|
|
||||||
|
This pattern is used by all tools (CRUD, file-system, etc.) via
|
||||||
|
``execute_on_client()`` in ``ws_context.py``.
|
||||||
|
|
||||||
|
The ``device_manager`` module-level singleton is imported by both the
|
||||||
|
device WS route and the agent runner.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
|
from fastapi import WebSocket
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DeviceConnection:
|
||||||
|
"""State for a single connected Electron device."""
|
||||||
|
|
||||||
|
ws: WebSocket
|
||||||
|
device_id: str
|
||||||
|
# Futures indexed by tool_call id — resolved when tool_result arrives.
|
||||||
|
pending_calls: dict[str, asyncio.Future[dict]] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
|
class DeviceConnectionManager:
|
||||||
|
"""Singleton registry of active Electron WebSocket connections.
|
||||||
|
|
||||||
|
Thread/task safety note: asyncio is single-threaded by design. All
|
||||||
|
mutations happen inside await-points on the main event loop, so no
|
||||||
|
locking is required for the in-memory dicts.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._connections: dict[str, DeviceConnection] = {}
|
||||||
|
|
||||||
|
# ── Registration ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def register(self, user_id: str, device_id: str, ws: WebSocket) -> None:
|
||||||
|
"""Store the active connection for *user_id*, replacing any previous one."""
|
||||||
|
if user_id in self._connections:
|
||||||
|
old = self._connections[user_id]
|
||||||
|
logger.info(
|
||||||
|
"device_manager: replacing existing connection for user=%s device=%s",
|
||||||
|
user_id,
|
||||||
|
old.device_id,
|
||||||
|
)
|
||||||
|
# Cancel any futures that were waiting on the old connection.
|
||||||
|
for fut in old.pending_calls.values():
|
||||||
|
if not fut.done():
|
||||||
|
fut.cancel()
|
||||||
|
self._connections[user_id] = DeviceConnection(ws=ws, device_id=device_id)
|
||||||
|
logger.info(
|
||||||
|
"device_manager: registered user=%s device=%s", user_id, device_id
|
||||||
|
)
|
||||||
|
|
||||||
|
def unregister(self, user_id: str) -> None:
|
||||||
|
"""Remove the connection for *user_id* and cancel any pending futures."""
|
||||||
|
conn = self._connections.pop(user_id, None)
|
||||||
|
if conn is None:
|
||||||
|
return
|
||||||
|
for fut in conn.pending_calls.values():
|
||||||
|
if not fut.done():
|
||||||
|
fut.cancel()
|
||||||
|
logger.info("device_manager: unregistered user=%s", user_id)
|
||||||
|
|
||||||
|
# ── Presence queries ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
def get_ws(self, user_id: str) -> WebSocket | None:
|
||||||
|
"""Return the active WebSocket for *user_id*, or ``None`` if offline."""
|
||||||
|
conn = self._connections.get(user_id)
|
||||||
|
return conn.ws if conn else None
|
||||||
|
|
||||||
|
def is_online(self, user_id: str, device_id: str | None = None) -> bool:
|
||||||
|
"""Return ``True`` if the user has an active connection.
|
||||||
|
|
||||||
|
If *device_id* is provided also checks that it matches the connected device.
|
||||||
|
"""
|
||||||
|
conn = self._connections.get(user_id)
|
||||||
|
if conn is None:
|
||||||
|
return False
|
||||||
|
if device_id is not None:
|
||||||
|
return conn.device_id == device_id
|
||||||
|
return True
|
||||||
|
|
||||||
|
# ── Frame sending ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def send_frame(self, user_id: str, frame: dict) -> None:
|
||||||
|
"""Send *frame* as a JSON text message to the device.
|
||||||
|
|
||||||
|
Raises ``RuntimeError`` if the user is not connected.
|
||||||
|
"""
|
||||||
|
conn = self._connections.get(user_id)
|
||||||
|
if conn is None:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"send_frame: user {user_id!r} is not connected"
|
||||||
|
)
|
||||||
|
await conn.ws.send_text(json.dumps(frame))
|
||||||
|
|
||||||
|
# ── Tool-call round-trip ──────────────────────────────────────────
|
||||||
|
|
||||||
|
def create_pending_call(
|
||||||
|
self, user_id: str, call_id: str
|
||||||
|
) -> asyncio.Future[dict]:
|
||||||
|
"""Register a Future that will be resolved when the tool_result arrives.
|
||||||
|
|
||||||
|
Raises ``RuntimeError`` if the user is not connected.
|
||||||
|
"""
|
||||||
|
conn = self._connections.get(user_id)
|
||||||
|
if conn is None:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"create_pending_call: user {user_id!r} is not connected"
|
||||||
|
)
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
fut: asyncio.Future[dict] = loop.create_future()
|
||||||
|
conn.pending_calls[call_id] = fut
|
||||||
|
return fut
|
||||||
|
|
||||||
|
def resolve_pending_call(
|
||||||
|
self, user_id: str, call_id: str, result: dict
|
||||||
|
) -> None:
|
||||||
|
"""Fulfil the Future registered under *call_id* with the Electron result.
|
||||||
|
|
||||||
|
No-ops if the call_id is unknown (already timed out or cancelled).
|
||||||
|
"""
|
||||||
|
conn = self._connections.get(user_id)
|
||||||
|
if conn is None:
|
||||||
|
return
|
||||||
|
fut = conn.pending_calls.pop(call_id, None)
|
||||||
|
if fut is not None and not fut.done():
|
||||||
|
fut.set_result(result)
|
||||||
|
|
||||||
|
|
||||||
|
# Module-level singleton — import this everywhere.
|
||||||
|
device_manager = DeviceConnectionManager()
|
||||||
34
app/core/embeddings.py
Normal file
34
app/core/embeddings.py
Normal file
@@ -0,0 +1,34 @@
|
|||||||
|
"""OpenAI embedding helper for associative memory tier.
|
||||||
|
|
||||||
|
Single public function: ``embed_text(text) -> list[float] | None``.
|
||||||
|
Returns None on any failure — callers must implement a keyword fallback.
|
||||||
|
Never raises; all exceptions are logged as warnings.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from openai import AsyncOpenAI
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_MAX_INPUT_CHARS = 8000
|
||||||
|
_EMBEDDING_MODEL = "text-embedding-3-small"
|
||||||
|
|
||||||
|
|
||||||
|
async def embed_text(text: str) -> list[float] | None:
|
||||||
|
"""Call OpenAI text-embedding-3-small. Return None on failure (caller falls back to keyword)."""
|
||||||
|
try:
|
||||||
|
client = AsyncOpenAI()
|
||||||
|
truncated = text[:_MAX_INPUT_CHARS]
|
||||||
|
response = await client.embeddings.create(
|
||||||
|
input=truncated,
|
||||||
|
model=_EMBEDDING_MODEL,
|
||||||
|
)
|
||||||
|
result: list[float] = response.data[0].embedding
|
||||||
|
logger.debug("embeddings: embed_text dims=%d", len(result))
|
||||||
|
return result
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("embeddings: embed_text failed: %s", exc)
|
||||||
|
return None
|
||||||
@@ -1,222 +0,0 @@
|
|||||||
"""Execution Plan generator — builder, template registry, and LRU plan cache."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from collections import OrderedDict
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
from app.schemas import ExecutionPlan, PlanStep
|
|
||||||
|
|
||||||
|
|
||||||
# ── Prompt Template Registry ──────────────────────────────────────────
|
|
||||||
|
|
||||||
|
|
||||||
class PromptTemplateRegistry:
|
|
||||||
"""Server-side store mapping template IDs to prompt text.
|
|
||||||
|
|
||||||
Clients only ever receive template IDs (e.g. ``"tpl_task_agent_default"``).
|
|
||||||
The actual prompt text is resolved here on the server, keeping prompt IP
|
|
||||||
out of API responses.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self) -> None:
|
|
||||||
self._templates: dict[str, str] = {}
|
|
||||||
|
|
||||||
def register(self, template_id: str, prompt_text: str) -> None:
|
|
||||||
self._templates[template_id] = prompt_text
|
|
||||||
|
|
||||||
def get(self, template_id: str) -> str:
|
|
||||||
"""Resolve a template ID to its prompt text.
|
|
||||||
|
|
||||||
Raises ``KeyError`` if the template is not registered.
|
|
||||||
"""
|
|
||||||
text = self._templates.get(template_id)
|
|
||||||
if text is None:
|
|
||||||
raise KeyError(f"Template not found: {template_id!r}")
|
|
||||||
return text
|
|
||||||
|
|
||||||
def has(self, template_id: str) -> bool:
|
|
||||||
return template_id in self._templates
|
|
||||||
|
|
||||||
def list_ids(self) -> list[str]:
|
|
||||||
"""Return all registered template IDs (never the text)."""
|
|
||||||
return list(self._templates.keys())
|
|
||||||
|
|
||||||
|
|
||||||
# ── Execution Plan Builder ────────────────────────────────────────────
|
|
||||||
|
|
||||||
|
|
||||||
class ExecutionPlanBuilder:
|
|
||||||
"""Fluent builder for ``ExecutionPlan`` objects.
|
|
||||||
|
|
||||||
Example::
|
|
||||||
|
|
||||||
plan = (
|
|
||||||
ExecutionPlanBuilder("task_agent")
|
|
||||||
.add_llm_step("tpl_task_agent_default", {"message": user_msg})
|
|
||||||
.add_data_step("create_record", data_from_step=0)
|
|
||||||
.build()
|
|
||||||
)
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, agent: str) -> None:
|
|
||||||
self._agent = agent
|
|
||||||
self._steps: list[PlanStep] = []
|
|
||||||
|
|
||||||
# ── step adders ──────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def add_step(
|
|
||||||
self, action: str, params: dict[str, Any] | None = None
|
|
||||||
) -> ExecutionPlanBuilder:
|
|
||||||
"""Append a generic action step with optional parameters."""
|
|
||||||
self._steps.append(PlanStep(action=action, variables=params))
|
|
||||||
return self
|
|
||||||
|
|
||||||
def add_llm_step(
|
|
||||||
self, template_id: str, variables: dict[str, Any] | None = None
|
|
||||||
) -> ExecutionPlanBuilder:
|
|
||||||
"""Append an LLM step referencing a server-side template by ID."""
|
|
||||||
self._steps.append(
|
|
||||||
PlanStep(action="llm", prompt_template=template_id, variables=variables)
|
|
||||||
)
|
|
||||||
return self
|
|
||||||
|
|
||||||
def add_data_step(self, action: str, data_from_step: int) -> ExecutionPlanBuilder:
|
|
||||||
"""Append a step whose input comes from the output of an earlier step."""
|
|
||||||
self._steps.append(PlanStep(action=action, data_from_step=data_from_step))
|
|
||||||
return self
|
|
||||||
|
|
||||||
# ── build ────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def build(self) -> ExecutionPlan:
|
|
||||||
"""Validate step references and return the ``ExecutionPlan``.
|
|
||||||
|
|
||||||
Raises ``ValueError`` if any ``data_from_step`` references a
|
|
||||||
non-existent or future step index.
|
|
||||||
"""
|
|
||||||
for i, step in enumerate(self._steps):
|
|
||||||
if step.data_from_step is not None:
|
|
||||||
if not (0 <= step.data_from_step < i):
|
|
||||||
raise ValueError(
|
|
||||||
f"Step {i}: data_from_step={step.data_from_step} must "
|
|
||||||
f"reference a preceding step index in range 0..{i - 1}"
|
|
||||||
)
|
|
||||||
return ExecutionPlan(agent=self._agent, steps=list(self._steps))
|
|
||||||
|
|
||||||
|
|
||||||
# ── Plan Cache (LRU) ──────────────────────────────────────────────────
|
|
||||||
|
|
||||||
|
|
||||||
class PlanCache:
|
|
||||||
"""In-memory LRU cache for ``ExecutionPlan`` objects.
|
|
||||||
|
|
||||||
Plans stored here are accessible as playbooks via ``get_all_playbooks()``.
|
|
||||||
The cache also serves as a runtime memoisation layer so that repeated
|
|
||||||
identical intent classifications can skip re-building the plan.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, maxsize: int = 1000) -> None:
|
|
||||||
self._maxsize = maxsize
|
|
||||||
self._cache: OrderedDict[str, ExecutionPlan] = OrderedDict()
|
|
||||||
|
|
||||||
def cache_plan(self, key: str, plan: ExecutionPlan) -> None:
|
|
||||||
"""Store *plan* under *key*, evicting the LRU entry if at capacity."""
|
|
||||||
if key in self._cache:
|
|
||||||
del self._cache[key] # remove so re-insertion places it at the end
|
|
||||||
elif len(self._cache) >= self._maxsize:
|
|
||||||
self._cache.popitem(last=False) # evict least-recently-used
|
|
||||||
self._cache[key] = plan
|
|
||||||
|
|
||||||
def get_plan(self, key: str) -> ExecutionPlan | None:
|
|
||||||
"""Return the cached plan for *key*, or ``None`` if not present.
|
|
||||||
|
|
||||||
Accessing a plan marks it as most-recently used.
|
|
||||||
"""
|
|
||||||
if key not in self._cache:
|
|
||||||
return None
|
|
||||||
self._cache.move_to_end(key)
|
|
||||||
return self._cache[key]
|
|
||||||
|
|
||||||
def get_all_playbooks(self) -> list[ExecutionPlan]:
|
|
||||||
"""Return all cached plans (most-recently used last)."""
|
|
||||||
return list(self._cache.values())
|
|
||||||
|
|
||||||
|
|
||||||
# ── Module-level singletons ───────────────────────────────────────────
|
|
||||||
|
|
||||||
template_registry = PromptTemplateRegistry()
|
|
||||||
plan_cache = PlanCache()
|
|
||||||
|
|
||||||
|
|
||||||
def _register_builtin_templates() -> None:
|
|
||||||
"""Register the built-in server-side prompt templates.
|
|
||||||
|
|
||||||
These strings never leave the server. Clients only receive the IDs.
|
|
||||||
"""
|
|
||||||
_tpls: dict[str, str] = {
|
|
||||||
"tpl_task_agent_default": (
|
|
||||||
"You are a task management assistant. Help the user create, update, "
|
|
||||||
"list, and track tasks. Use correct status values (todo, in_progress, "
|
|
||||||
"done) and priority values (high, medium, low) from the workspace model."
|
|
||||||
),
|
|
||||||
"tpl_checkpoint_agent_default": (
|
|
||||||
"You are a project checkpoint assistant. Help the user create and manage "
|
|
||||||
"milestone checkpoints on their projects. Every checkpoint requires a "
|
|
||||||
"project_id and a date expressed as a Unix timestamp in milliseconds."
|
|
||||||
),
|
|
||||||
"tpl_project_agent_default": (
|
|
||||||
"You are a project management assistant. Help the user create, find, "
|
|
||||||
"update, and archive projects. Projects have a name, an optional client, "
|
|
||||||
"and a status of either active or archived."
|
|
||||||
),
|
|
||||||
"tpl_note_agent_default": (
|
|
||||||
"You are a note-taking assistant. Help the user create, retrieve, update, "
|
|
||||||
"and delete Markdown notes. Notes can optionally be linked to a project."
|
|
||||||
),
|
|
||||||
"tpl_task_extract_from_project": (
|
|
||||||
"Extract all actionable tasks from the provided project context. "
|
|
||||||
"Return a structured list of tasks, each with a title, inferred priority "
|
|
||||||
"(high, medium, or low), suggested status (todo), and a due_date in "
|
|
||||||
"milliseconds where a deadline can be inferred."
|
|
||||||
),
|
|
||||||
"tpl_note_weekly_summary": (
|
|
||||||
"Generate a weekly project summary note from the provided workspace data. "
|
|
||||||
"Include: tasks completed this week, tasks due soon, active projects, "
|
|
||||||
"and upcoming checkpoints. Format the output as clean Markdown."
|
|
||||||
),
|
|
||||||
}
|
|
||||||
for tid, text in _tpls.items():
|
|
||||||
template_registry.register(tid, text)
|
|
||||||
|
|
||||||
|
|
||||||
def _load_playbooks() -> None:
|
|
||||||
"""Pre-build and cache the built-in playbooks."""
|
|
||||||
playbooks: list[tuple[str, ExecutionPlan]] = [
|
|
||||||
(
|
|
||||||
"create_tasks_from_project",
|
|
||||||
ExecutionPlanBuilder("project_agent")
|
|
||||||
.add_llm_step(
|
|
||||||
"tpl_task_extract_from_project",
|
|
||||||
{"source": "project_context"},
|
|
||||||
)
|
|
||||||
.add_data_step("create_record", data_from_step=0)
|
|
||||||
.build(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"generate_weekly_note",
|
|
||||||
ExecutionPlanBuilder("note_agent")
|
|
||||||
.add_llm_step(
|
|
||||||
"tpl_note_weekly_summary",
|
|
||||||
{"period": "last_7_days"},
|
|
||||||
)
|
|
||||||
.add_data_step("create_record", data_from_step=0)
|
|
||||||
.build(),
|
|
||||||
),
|
|
||||||
]
|
|
||||||
for key, plan in playbooks:
|
|
||||||
plan_cache.cache_plan(key, plan)
|
|
||||||
|
|
||||||
|
|
||||||
# Initialise on module load
|
|
||||||
_register_builtin_templates()
|
|
||||||
_load_playbooks()
|
|
||||||
183
app/core/folder_indexer.py
Normal file
183
app/core/folder_indexer.py
Normal file
@@ -0,0 +1,183 @@
|
|||||||
|
"""Per-file summarisation for project folder integration."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import io
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from langchain_core.messages import HumanMessage, SystemMessage
|
||||||
|
from pypdf import PdfReader
|
||||||
|
from docx import Document as DocxDocument
|
||||||
|
|
||||||
|
from app.core.langfuse_client import (
|
||||||
|
compile_prompt,
|
||||||
|
extract_usage,
|
||||||
|
get_langfuse,
|
||||||
|
get_prompt_or_fallback,
|
||||||
|
)
|
||||||
|
from app.core.llm import get_llm
|
||||||
|
|
||||||
|
_TEXT_FALLBACK = (
|
||||||
|
"You are summarising a file for an AI assistant that helps the user manage a project.\n"
|
||||||
|
"Produce a single sentence (<=30 words, <=200 chars) that captures the file's purpose "
|
||||||
|
"and most important detail.\nFile extension: {ext}\nFile name: {name}\nContent (truncated if long):\n{content}"
|
||||||
|
)
|
||||||
|
_IMAGE_FALLBACK = (
|
||||||
|
"You are summarising an image attached to a project folder.\n"
|
||||||
|
"Produce a single sentence (<=30 words, <=200 chars) describing what the image shows "
|
||||||
|
"and any obvious purpose (logo, screenshot, diagram, photo of a whiteboard, etc.)."
|
||||||
|
)
|
||||||
|
_MAX_INPUT_CHARS = 6000
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class IndexResult:
|
||||||
|
summary: str
|
||||||
|
tokens_used: int
|
||||||
|
|
||||||
|
|
||||||
|
async def _llm_text(messages: list) -> object:
|
||||||
|
"""Make the LLM call for text summarisation.
|
||||||
|
|
||||||
|
Defined as a standalone async function so tests can patch it cleanly
|
||||||
|
without needing to mock the LLM object itself.
|
||||||
|
"""
|
||||||
|
llm = get_llm(model="gpt-4o-mini", temperature=0.2)
|
||||||
|
return await llm.ainvoke(messages)
|
||||||
|
|
||||||
|
|
||||||
|
async def _llm_vision(messages: list) -> object:
|
||||||
|
"""Make the LLM call for vision (image) summarisation.
|
||||||
|
|
||||||
|
Accepts the message list and returns the response directly, mirroring
|
||||||
|
the ``_llm_text`` caller pattern so tests can patch it at the module level.
|
||||||
|
"""
|
||||||
|
llm = get_llm(model="gpt-4o-mini", temperature=0.2)
|
||||||
|
return await llm.ainvoke(messages)
|
||||||
|
|
||||||
|
|
||||||
|
async def summarize_image(*, image_b64: str, mime: str, file_name: str | None = None) -> IndexResult:
|
||||||
|
"""Return a compact summary of an image file using vision.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
image_b64:
|
||||||
|
Base64-encoded image bytes.
|
||||||
|
mime:
|
||||||
|
MIME type of the image, e.g. ``"image/png"``.
|
||||||
|
file_name:
|
||||||
|
Optional file name, attached to the Langfuse trace as input metadata.
|
||||||
|
"""
|
||||||
|
template, prompt_obj = get_prompt_or_fallback("folder_file_summary_image", _IMAGE_FALLBACK)
|
||||||
|
messages = [
|
||||||
|
SystemMessage(content=template),
|
||||||
|
HumanMessage(content=[
|
||||||
|
{"type": "text", "text": "Summarise this image."},
|
||||||
|
{"type": "image_url", "image_url": {"url": f"data:{mime};base64,{image_b64}"}},
|
||||||
|
]),
|
||||||
|
]
|
||||||
|
lf = get_langfuse()
|
||||||
|
if lf is not None:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="folder-summarize-image",
|
||||||
|
model="gpt-4o-mini",
|
||||||
|
prompt=prompt_obj,
|
||||||
|
input={"file_name": file_name, "mime": mime},
|
||||||
|
) as gen:
|
||||||
|
response = await _llm_vision(messages)
|
||||||
|
usage = extract_usage(response)
|
||||||
|
gen.update(output=response.content, usage_details=usage)
|
||||||
|
else:
|
||||||
|
response = await _llm_vision(messages)
|
||||||
|
usage = extract_usage(response)
|
||||||
|
summary = (response.content or "").strip()[:500]
|
||||||
|
return IndexResult(summary=summary, tokens_used=usage.get("total", 0))
|
||||||
|
|
||||||
|
|
||||||
|
async def summarize_text(*, content: str, ext: str, name: str) -> IndexResult:
|
||||||
|
"""Return a compact summary of a text file.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
content:
|
||||||
|
Raw text content of the file (will be truncated to _MAX_INPUT_CHARS).
|
||||||
|
ext:
|
||||||
|
File extension including the leading dot, e.g. ``".md"``.
|
||||||
|
name:
|
||||||
|
File name, e.g. ``"kickoff.md"``.
|
||||||
|
"""
|
||||||
|
template, prompt_obj = get_prompt_or_fallback("folder_file_summary_text", _TEXT_FALLBACK)
|
||||||
|
truncated = content[:_MAX_INPUT_CHARS]
|
||||||
|
compiled = compile_prompt(template, prompt_obj, ext=ext, name=name, content=truncated)
|
||||||
|
messages = [
|
||||||
|
SystemMessage(content=compiled),
|
||||||
|
HumanMessage(content="Summarise this file."),
|
||||||
|
]
|
||||||
|
lf = get_langfuse()
|
||||||
|
if lf is not None:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="folder-summarize-text",
|
||||||
|
model="gpt-4o-mini",
|
||||||
|
prompt=prompt_obj,
|
||||||
|
input={"file_name": name, "ext": ext, "content_chars": len(truncated)},
|
||||||
|
) as gen:
|
||||||
|
response = await _llm_text(messages)
|
||||||
|
usage = extract_usage(response)
|
||||||
|
gen.update(output=response.content, usage_details=usage)
|
||||||
|
else:
|
||||||
|
response = await _llm_text(messages)
|
||||||
|
usage = extract_usage(response)
|
||||||
|
summary = (response.content or "").strip()[:500]
|
||||||
|
return IndexResult(summary=summary, tokens_used=usage.get("total", 0))
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_pdf_text(pdf_b64: str) -> str:
|
||||||
|
buf = io.BytesIO(base64.b64decode(pdf_b64))
|
||||||
|
reader = PdfReader(buf)
|
||||||
|
parts: list[str] = []
|
||||||
|
for page in reader.pages:
|
||||||
|
try:
|
||||||
|
parts.append(page.extract_text() or "")
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
return "\n".join(parts).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_docx_text(docx_b64: str) -> str:
|
||||||
|
buf = io.BytesIO(base64.b64decode(docx_b64))
|
||||||
|
doc = DocxDocument(buf)
|
||||||
|
return "\n".join(p.text for p in doc.paragraphs if p.text).strip()
|
||||||
|
|
||||||
|
|
||||||
|
async def summarize_pdf(*, pdf_b64: str, name: str) -> IndexResult:
|
||||||
|
"""Return a compact summary of a PDF file.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
pdf_b64:
|
||||||
|
Base64-encoded PDF bytes.
|
||||||
|
name:
|
||||||
|
File name, e.g. ``"report.pdf"``.
|
||||||
|
"""
|
||||||
|
text = _extract_pdf_text(pdf_b64)
|
||||||
|
if not text:
|
||||||
|
return IndexResult(summary="Could not extract text", tokens_used=0)
|
||||||
|
return await summarize_text(content=text, ext=".pdf", name=name)
|
||||||
|
|
||||||
|
|
||||||
|
async def summarize_docx(*, docx_b64: str, name: str) -> IndexResult:
|
||||||
|
"""Return a compact summary of a DOCX file.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
docx_b64:
|
||||||
|
Base64-encoded DOCX bytes.
|
||||||
|
name:
|
||||||
|
File name, e.g. ``"spec.docx"``.
|
||||||
|
"""
|
||||||
|
text = _extract_docx_text(docx_b64)
|
||||||
|
if not text:
|
||||||
|
return IndexResult(summary="Could not extract text", tokens_used=0)
|
||||||
|
return await summarize_text(content=text, ext=".docx", name=name)
|
||||||
190
app/core/langfuse_client.py
Normal file
190
app/core/langfuse_client.py
Normal file
@@ -0,0 +1,190 @@
|
|||||||
|
"""Langfuse observability — singleton client and prompt helpers.
|
||||||
|
|
||||||
|
If LANGFUSE_SECRET_KEY / LANGFUSE_PUBLIC_KEY are not set,
|
||||||
|
all helpers are no-ops so the app works without Langfuse configured.
|
||||||
|
|
||||||
|
Usage
|
||||||
|
-----
|
||||||
|
Tracing::
|
||||||
|
|
||||||
|
from app.core.langfuse_client import get_langfuse
|
||||||
|
|
||||||
|
lf = get_langfuse()
|
||||||
|
if lf:
|
||||||
|
with lf.start_as_current_observation(as_type="span", name="my-agent") as span:
|
||||||
|
span.update(input=user_message)
|
||||||
|
# ... do work ...
|
||||||
|
span.update(output=result)
|
||||||
|
lf.flush()
|
||||||
|
|
||||||
|
Prompt management::
|
||||||
|
|
||||||
|
from app.core.langfuse_client import get_prompt_or_fallback
|
||||||
|
|
||||||
|
text, prompt_obj = get_prompt_or_fallback("home_system", FALLBACK_PROMPT)
|
||||||
|
# Use text as the system prompt; pass prompt_obj to generations for linking.
|
||||||
|
|
||||||
|
Linking a prompt to a generation::
|
||||||
|
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="llm-call",
|
||||||
|
model="gpt-4o",
|
||||||
|
prompt=prompt_obj, # links generation → prompt version in the UI
|
||||||
|
input=messages,
|
||||||
|
) as gen:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
gen.update(output=response.content, usage=_usage(response))
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
import logging
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from typing import Any, Generator
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_client: Any = None
|
||||||
|
_initialized: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
def get_langfuse() -> Any | None:
|
||||||
|
"""Return the Langfuse singleton, or ``None`` when not configured."""
|
||||||
|
global _client, _initialized
|
||||||
|
if _initialized:
|
||||||
|
return _client
|
||||||
|
_initialized = True
|
||||||
|
|
||||||
|
from app.config.settings import settings # local import to avoid circular deps
|
||||||
|
|
||||||
|
if not settings.LANGFUSE_SECRET_KEY or not settings.LANGFUSE_PUBLIC_KEY:
|
||||||
|
logger.debug("langfuse: not configured — observability disabled")
|
||||||
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
|
from langfuse import Langfuse
|
||||||
|
|
||||||
|
_client = Langfuse(
|
||||||
|
secret_key=settings.LANGFUSE_SECRET_KEY,
|
||||||
|
public_key=settings.LANGFUSE_PUBLIC_KEY,
|
||||||
|
host=settings.LANGFUSE_BASE_URL,
|
||||||
|
)
|
||||||
|
logger.info("langfuse: client initialized host=%s", settings.LANGFUSE_BASE_URL)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("langfuse: failed to initialize: %s", exc)
|
||||||
|
_client = None
|
||||||
|
|
||||||
|
return _client
|
||||||
|
|
||||||
|
|
||||||
|
def get_prompt_or_fallback(name: str, fallback: str) -> tuple[str, Any]:
|
||||||
|
"""Fetch a text prompt from Langfuse; fall back to ``fallback`` on any error.
|
||||||
|
|
||||||
|
Returns ``(raw_template, prompt_obj_or_None)``.
|
||||||
|
|
||||||
|
* ``raw_template`` — the uncompiled template string. Do NOT call ``.format()``
|
||||||
|
on it directly; use :func:`compile_prompt` instead so the correct variable
|
||||||
|
syntax is applied (``{{var}}`` for Langfuse, ``{var}`` for the fallback).
|
||||||
|
* ``prompt_obj`` — the Langfuse prompt object, or ``None`` when Langfuse is
|
||||||
|
unavailable / the fetch failed. Pass this to generation observations so
|
||||||
|
Langfuse links the generation to the exact prompt version in the UI.
|
||||||
|
"""
|
||||||
|
lf = get_langfuse()
|
||||||
|
if lf is None:
|
||||||
|
return fallback, None
|
||||||
|
|
||||||
|
try:
|
||||||
|
prompt = lf.get_prompt(name, label="production", fallback=fallback)
|
||||||
|
# For text-type prompts .prompt holds the raw template string.
|
||||||
|
raw = prompt.prompt if hasattr(prompt, "prompt") and isinstance(prompt.prompt, str) else fallback
|
||||||
|
return raw, prompt
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("langfuse: get_prompt %r failed: %s — using fallback", name, exc)
|
||||||
|
return fallback, None
|
||||||
|
|
||||||
|
|
||||||
|
def compile_prompt(template: str, prompt_obj: Any, **variables: Any) -> str:
|
||||||
|
"""Compile *template* with *variables*, choosing the right syntax.
|
||||||
|
|
||||||
|
* When *prompt_obj* is a real Langfuse prompt object, calls
|
||||||
|
``prompt_obj.compile(**variables)`` which handles ``{{variable}}``
|
||||||
|
substitution as defined in the Langfuse UI.
|
||||||
|
* When *prompt_obj* is ``None`` (Langfuse unavailable or fetch failed),
|
||||||
|
falls back to ``template.format(**variables)`` which handles the
|
||||||
|
``{variable}`` syntax used in the hardcoded fallback strings.
|
||||||
|
|
||||||
|
This keeps callers oblivious to which syntax is in use.
|
||||||
|
"""
|
||||||
|
if prompt_obj is not None:
|
||||||
|
try:
|
||||||
|
compiled = prompt_obj.compile(**variables)
|
||||||
|
# compile() returns a string for text prompts.
|
||||||
|
if isinstance(compiled, str):
|
||||||
|
return compiled
|
||||||
|
# Chat prompts return a list of dicts — join text parts.
|
||||||
|
if isinstance(compiled, list):
|
||||||
|
return "\n".join(
|
||||||
|
m.get("content", "") for m in compiled if isinstance(m, dict)
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"langfuse: compile failed for prompt %r: %s — falling back to .format()",
|
||||||
|
getattr(prompt_obj, "name", "?"),
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
return template.format(**variables)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_usage(response: Any) -> dict[str, int]:
|
||||||
|
"""Extract token usage from a LangChain AI message into Langfuse format."""
|
||||||
|
meta = getattr(response, "usage_metadata", None)
|
||||||
|
if not meta:
|
||||||
|
return {}
|
||||||
|
return {
|
||||||
|
"input": int(meta.get("input_tokens", 0)),
|
||||||
|
"output": int(meta.get("output_tokens", 0)),
|
||||||
|
"total": int(meta.get("total_tokens", 0)),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def hash_user_id(user_id: str) -> str:
|
||||||
|
"""Return a SHA-256 hash of *user_id* for use as Langfuse ``user_id``.
|
||||||
|
|
||||||
|
This avoids sending raw database UUIDs to external observability services
|
||||||
|
while still providing a stable, deterministic identifier for per-user
|
||||||
|
metrics in the Langfuse dashboard.
|
||||||
|
"""
|
||||||
|
return hashlib.sha256(user_id.encode()).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def langfuse_context(
|
||||||
|
user_id: str | None = None,
|
||||||
|
session_id: str | None = None,
|
||||||
|
) -> Generator[None, None, None]:
|
||||||
|
"""Propagate ``user_id`` (hashed) and ``session_id`` to all Langfuse observations.
|
||||||
|
|
||||||
|
No-op when Langfuse is not configured or parameters are empty.
|
||||||
|
"""
|
||||||
|
lf = get_langfuse()
|
||||||
|
if lf is None or (not user_id and not session_id):
|
||||||
|
yield
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
from langfuse import propagate_attributes
|
||||||
|
except ImportError:
|
||||||
|
logger.debug("langfuse: propagate_attributes not available — skipping context")
|
||||||
|
yield
|
||||||
|
return
|
||||||
|
|
||||||
|
attrs: dict[str, str] = {}
|
||||||
|
if user_id:
|
||||||
|
attrs["user_id"] = hash_user_id(user_id)
|
||||||
|
if session_id:
|
||||||
|
attrs["session_id"] = session_id
|
||||||
|
|
||||||
|
with propagate_attributes(**attrs):
|
||||||
|
yield
|
||||||
101
app/core/llm.py
101
app/core/llm.py
@@ -1,6 +1,6 @@
|
|||||||
"""LLM factory — centralised model instantiation via LiteLLM.
|
"""LLM factory — centralised model instantiation via LiteLLM.
|
||||||
|
|
||||||
Every agent and the orchestrator call ``get_llm()`` or ``get_router_llm()``
|
Every agent and the orchestrator call ``get_llm()``
|
||||||
instead of directly constructing a provider-specific class. The model string
|
instead of directly constructing a provider-specific class. The model string
|
||||||
follows the `LiteLLM model naming convention
|
follows the `LiteLLM model naming convention
|
||||||
<https://docs.litellm.ai/docs/providers>`_:
|
<https://docs.litellm.ai/docs/providers>`_:
|
||||||
@@ -11,19 +11,37 @@ follows the `LiteLLM model naming convention
|
|||||||
* Ollama: ``ollama/llama3``
|
* Ollama: ``ollama/llama3``
|
||||||
* Bedrock: ``bedrock/anthropic.claude-v2``
|
* Bedrock: ``bedrock/anthropic.claude-v2``
|
||||||
|
|
||||||
Switch providers by changing **LLM_MODEL** / **LLM_ROUTER_MODEL** in ``.env``
|
Switch providers by changing **LLM_MODEL** in ``.env``
|
||||||
— no code changes required.
|
— no code changes required.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import warnings
|
||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
from openai import AsyncOpenAI
|
from openai import AsyncOpenAI
|
||||||
|
import litellm
|
||||||
|
|
||||||
from langchain_openai import ChatOpenAI
|
from langchain_openai import ChatOpenAI
|
||||||
|
from langchain_litellm import ChatLiteLLM
|
||||||
from litellm import get_supported_openai_params # noqa: F401 – validates install
|
from litellm import get_supported_openai_params # noqa: F401 – validates install
|
||||||
|
|
||||||
from app.config.settings import settings
|
from app.config.settings import settings
|
||||||
|
|
||||||
|
# Some models (e.g. gpt-5, o-series) reject unsupported params like temperature.
|
||||||
|
# Drop them silently instead of raising UnsupportedParamsError.
|
||||||
|
litellm.drop_params = True
|
||||||
|
|
||||||
|
# Some provider responses include a plain dict in the `usage` field where a
|
||||||
|
# richer Pydantic model is expected. This warning is noisy but non-fatal.
|
||||||
|
warnings.filterwarnings(
|
||||||
|
"ignore",
|
||||||
|
message=r"PydanticSerializationUnexpectedValue\(Expected `ResponseAPIUsage`",
|
||||||
|
category=UserWarning,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _api_key_for_model(model: str) -> str | None:
|
def _api_key_for_model(model: str) -> str | None:
|
||||||
"""Return the most appropriate API key for the given LiteLLM model string."""
|
"""Return the most appropriate API key for the given LiteLLM model string."""
|
||||||
@@ -31,6 +49,16 @@ def _api_key_for_model(model: str) -> str | None:
|
|||||||
return settings.ANTHROPIC_API_KEY or None
|
return settings.ANTHROPIC_API_KEY or None
|
||||||
if model.startswith("gemini/") or model.startswith("google/"):
|
if model.startswith("gemini/") or model.startswith("google/"):
|
||||||
return settings.GOOGLE_API_KEY or None
|
return settings.GOOGLE_API_KEY or None
|
||||||
|
if model.startswith("cerebras/"):
|
||||||
|
return settings.CEREBRAS_API_KEY or None
|
||||||
|
if model.startswith("groq/"):
|
||||||
|
return settings.GROQ_API_KEY or None
|
||||||
|
if model.startswith("deepseek/"):
|
||||||
|
return settings.DEEPSEEK_API_KEY or None
|
||||||
|
if model.startswith("github_copilot/"):
|
||||||
|
# GitHub Copilot uses OAuth device-flow tokens managed by LiteLLM.
|
||||||
|
# No API key is required; returning None lets LiteLLM handle auth.
|
||||||
|
return None
|
||||||
# Default: OpenAI-compatible (covers plain model names like "gpt-4o")
|
# Default: OpenAI-compatible (covers plain model names like "gpt-4o")
|
||||||
return settings.OPENAI_API_KEY or None
|
return settings.OPENAI_API_KEY or None
|
||||||
|
|
||||||
@@ -39,7 +67,7 @@ def get_llm(
|
|||||||
*,
|
*,
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
temperature: float = 0,
|
temperature: float = 0,
|
||||||
) -> ChatOpenAI:
|
) -> ChatOpenAI | ChatLiteLLM:
|
||||||
"""Return a LangChain chat model backed by LiteLLM.
|
"""Return a LangChain chat model backed by LiteLLM.
|
||||||
|
|
||||||
LiteLLM exposes an OpenAI-compatible API, so we use ``ChatOpenAI`` pointed
|
LiteLLM exposes an OpenAI-compatible API, so we use ``ChatOpenAI`` pointed
|
||||||
@@ -55,6 +83,16 @@ def get_llm(
|
|||||||
Sampling temperature. ``0`` = deterministic.
|
Sampling temperature. ``0`` = deterministic.
|
||||||
"""
|
"""
|
||||||
model = model or settings.LLM_MODEL
|
model = model or settings.LLM_MODEL
|
||||||
|
|
||||||
|
# Point LiteLLM to the custom token directory when configured.
|
||||||
|
if settings.GITHUB_COPILOT_TOKEN_DIR:
|
||||||
|
os.environ.setdefault("GITHUB_COPILOT_TOKEN_DIR", settings.GITHUB_COPILOT_TOKEN_DIR)
|
||||||
|
|
||||||
|
# Use ChatLiteLLM for provider-prefixed models (github_copilot/, anthropic/, etc.)
|
||||||
|
# so LiteLLM handles routing and auth. ChatOpenAI for plain OpenAI model names.
|
||||||
|
if "/" in model:
|
||||||
|
return ChatLiteLLM(model=model, temperature=temperature)
|
||||||
|
|
||||||
return ChatOpenAI(
|
return ChatOpenAI(
|
||||||
model=model,
|
model=model,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
@@ -62,19 +100,58 @@ def get_llm(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def get_router_llm(
|
_AGENT_MODEL_SETTINGS: dict[str, Callable[[], str]] = {
|
||||||
|
"classifier": lambda: settings.LLM_MODEL_CLASSIFIER or settings.LLM_MODEL,
|
||||||
|
"home-agent": lambda: settings.LLM_MODEL_HOME_AGENT or settings.LLM_MODEL,
|
||||||
|
"floating-agent": lambda: settings.LLM_MODEL_FLOATING_AGENT or settings.LLM_MODEL,
|
||||||
|
"unified-processor": lambda: settings.LLM_MODEL_UNIFIED_PROCESSOR or settings.LLM_MODEL,
|
||||||
|
"cloud-processor": lambda: settings.LLM_MODEL_CLOUD_PROCESSOR or settings.LLM_MODEL,
|
||||||
|
"brief-agent": lambda: settings.LLM_MODEL_BRIEF_AGENT or settings.LLM_MODEL,
|
||||||
|
"task-brief-agent": lambda: settings.LLM_MODEL_TASK_BRIEF_AGENT or settings.LLM_MODEL,
|
||||||
|
"setup": lambda: settings.LLM_MODEL_SETUP_AGENT or settings.LLM_MODEL,
|
||||||
|
"memory-extractor": lambda: settings.LLM_MODEL_MEMORY_EXTRACTOR or "gpt-4o-mini",
|
||||||
|
"memory-miner": lambda: settings.LLM_MODEL_MEMORY_MINER or "gpt-4o-mini",
|
||||||
|
"memory-auditor": lambda: settings.LLM_MODEL_MEMORY_AUDITOR or settings.LLM_MODEL,
|
||||||
|
"note-summarizer": lambda: "gpt-4o-mini",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def model_for_agent(agent_name: str) -> str:
|
||||||
|
"""Return the resolved model string for *agent_name* (for Langfuse tracking)."""
|
||||||
|
return _AGENT_MODEL_SETTINGS.get(agent_name, lambda: settings.LLM_MODEL)()
|
||||||
|
|
||||||
|
|
||||||
|
def get_agent_llm(
|
||||||
|
agent_name: str,
|
||||||
*,
|
*,
|
||||||
temperature: float = 0,
|
temperature: float = 0,
|
||||||
) -> ChatOpenAI:
|
) -> ChatOpenAI | ChatLiteLLM:
|
||||||
"""Return the lighter model used for intent classification / routing."""
|
"""Return an LLM configured for *agent_name*, respecting per-agent overrides.
|
||||||
return get_llm(model=settings.LLM_ROUTER_MODEL, temperature=temperature)
|
|
||||||
|
Falls back to ``settings.LLM_MODEL`` for unknown agent names or when the
|
||||||
|
per-agent override is left empty in ``.env``.
|
||||||
|
"""
|
||||||
|
model = model_for_agent(agent_name)
|
||||||
|
return get_llm(model=model, temperature=temperature)
|
||||||
|
|
||||||
|
|
||||||
async def embed(text: str) -> list[float]:
|
async def embed(text: str) -> list[float]:
|
||||||
"""Return a 1536-dim embedding vector for *text* using text-embedding-3-small."""
|
"""Return an embedding vector for *text*.
|
||||||
|
|
||||||
|
Uses ``settings.LLM_EMBED_MODEL`` so the same provider switch in ``.env``
|
||||||
|
(e.g. ``github_copilot/text-embedding-3-small``) applies here without any
|
||||||
|
code changes. Falls back to the raw AsyncOpenAI client for plain OpenAI
|
||||||
|
model names to preserve existing behaviour.
|
||||||
|
"""
|
||||||
|
model = settings.LLM_EMBED_MODEL
|
||||||
|
|
||||||
|
if model.startswith("github_copilot/") or "/" in model:
|
||||||
|
# Use LiteLLM for all provider-prefixed models (Copilot, Bedrock, etc.)
|
||||||
|
# so the provider's auth mechanism is applied correctly.
|
||||||
|
response = await litellm.aembedding(model=model, input=[text])
|
||||||
|
return response.data[0]["embedding"]
|
||||||
|
|
||||||
|
# Plain OpenAI model name — use the raw AsyncOpenAI client (existing path).
|
||||||
client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY)
|
client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY)
|
||||||
response = await client.embeddings.create(
|
response = await client.embeddings.create(model=model, input=text)
|
||||||
model="text-embedding-3-small",
|
|
||||||
input=text,
|
|
||||||
)
|
|
||||||
return response.data[0].embedding
|
return response.data[0].embedding
|
||||||
|
|||||||
450
app/core/memory_extraction.py
Normal file
450
app/core/memory_extraction.py
Normal file
@@ -0,0 +1,450 @@
|
|||||||
|
"""Mem0-style Extract/Update pipeline — Phase 2.
|
||||||
|
|
||||||
|
Runs after every ``store_episode`` call to distil durable facts, preferences,
|
||||||
|
routines, and relations from the latest conversation turn.
|
||||||
|
|
||||||
|
Entry point: ``run_extraction(db, user_id, last_user_msg, last_assistant_msg, session_id)``
|
||||||
|
|
||||||
|
Design notes
|
||||||
|
------------
|
||||||
|
- Two gpt-4o-mini calls per turn: extract candidates, then decide action per candidate.
|
||||||
|
- Short-circuit: if no existing neighbours → ADD without a second LLM call (cost saving).
|
||||||
|
- Zero-trust: never logs decrypted user content; relation subject/object labels are
|
||||||
|
treated as identifiers (safe to log per spec).
|
||||||
|
- Must not raise into the request path — caller wraps in asyncio.create_task().
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from typing import Any, Literal
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.core.langfuse_client import get_langfuse, get_prompt_or_fallback, extract_usage, langfuse_context
|
||||||
|
from app.core.llm import get_agent_llm, model_for_agent
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ── Fallback prompts (used when Langfuse unavailable) ─────────────────────────
|
||||||
|
|
||||||
|
_EXTRACTION_FALLBACK = (
|
||||||
|
"You are a memory extractor for a personal AI secretary. Given the last conversation "
|
||||||
|
"turn, the user's core memory, and recent episode summaries, identify durable facts, "
|
||||||
|
"preferences, routines, and person/project relations worth remembering.\n\n"
|
||||||
|
"Output JSON matching this schema exactly:\n"
|
||||||
|
'{{"candidates": [{{"type": "<fact|preference|relation|routine>", '
|
||||||
|
'"content": "<short canonical statement>", '
|
||||||
|
'"target_tier": "<core|associative|relational|proactive>", '
|
||||||
|
'"subject": null, "predicate": null, "object": null, "confidence": 0.7}}]}}\n\n'
|
||||||
|
"Rules:\n"
|
||||||
|
"- Skip small talk, greetings, one-off questions.\n"
|
||||||
|
"- Max 5 candidates per call.\n"
|
||||||
|
"- Only extract durable information (still true next week).\n"
|
||||||
|
"- For type=relation: subject/predicate/object required.\n"
|
||||||
|
"- Default confidence=0.7.\n\n"
|
||||||
|
"## Last turn\n{last_turn}\n\n"
|
||||||
|
"## Core memory (current)\n{core_memory}\n\n"
|
||||||
|
"## Recent episodes\n{recent_episodes}"
|
||||||
|
)
|
||||||
|
|
||||||
|
_DECIDE_FALLBACK = (
|
||||||
|
"You are a memory update decision engine. Given a new memory candidate and a list of "
|
||||||
|
"existing memories from the same tier, decide what action to take.\n\n"
|
||||||
|
"Respond with exactly one word: ADD, UPDATE, DELETE, or NOOP.\n\n"
|
||||||
|
"- ADD: new information not in existing memories.\n"
|
||||||
|
"- UPDATE: contradicts or supersedes an existing memory.\n"
|
||||||
|
"- DELETE: states something is no longer true.\n"
|
||||||
|
"- NOOP: already captured accurately.\n\n"
|
||||||
|
"## New candidate\n{candidate}\n\n"
|
||||||
|
"## Existing memories (same tier, top neighbours)\n{existing_memories}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Pydantic schemas ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
class MemoryCandidate(BaseModel):
|
||||||
|
type: Literal["fact", "preference", "relation", "routine"]
|
||||||
|
content: str
|
||||||
|
target_tier: Literal["core", "associative", "relational", "proactive"]
|
||||||
|
subject: str | None = None
|
||||||
|
predicate: str | None = None
|
||||||
|
object: str | None = None
|
||||||
|
confidence: float = Field(default=0.7, ge=0.0, le=1.0)
|
||||||
|
|
||||||
|
|
||||||
|
class ExtractionResult(BaseModel):
|
||||||
|
candidates: list[MemoryCandidate] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Task 2.1 — Extract candidates ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def extract_candidates(
|
||||||
|
last_turn: str,
|
||||||
|
core_memory: dict[str, str],
|
||||||
|
recent_episodes: list[str],
|
||||||
|
) -> ExtractionResult:
|
||||||
|
"""Call gpt-4o-mini to extract memory candidates from the latest turn.
|
||||||
|
|
||||||
|
Returns an ExtractionResult (may be empty on failure — never raises).
|
||||||
|
"""
|
||||||
|
core_str = "\n".join(f"{k}: {v}" for k, v in core_memory.items()) or "(empty)"
|
||||||
|
episodes_str = "\n---\n".join(recent_episodes[-5:]) or "(none)"
|
||||||
|
|
||||||
|
template, prompt_obj = get_prompt_or_fallback("memory_extraction", _EXTRACTION_FALLBACK)
|
||||||
|
|
||||||
|
# Compile with Langfuse variable syntax ({{var}}) or fallback {var}
|
||||||
|
if prompt_obj is not None:
|
||||||
|
try:
|
||||||
|
system_text = prompt_obj.compile(
|
||||||
|
last_turn=last_turn,
|
||||||
|
core_memory=core_str,
|
||||||
|
recent_episodes=episodes_str,
|
||||||
|
)
|
||||||
|
if isinstance(system_text, list):
|
||||||
|
system_text = "\n".join(m.get("content", "") for m in system_text if isinstance(m, dict))
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_extraction: compile failed: %s", exc)
|
||||||
|
system_text = template.format(
|
||||||
|
last_turn=last_turn,
|
||||||
|
core_memory=core_str,
|
||||||
|
recent_episodes=episodes_str,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
system_text = template.format(
|
||||||
|
last_turn=last_turn,
|
||||||
|
core_memory=core_str,
|
||||||
|
recent_episodes=episodes_str,
|
||||||
|
)
|
||||||
|
|
||||||
|
llm = get_agent_llm("memory-extractor", temperature=0)
|
||||||
|
# Bind JSON mode so the model always returns parseable output.
|
||||||
|
llm_json = llm.bind(response_format={"type": "json_object"}) # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
lf = get_langfuse()
|
||||||
|
try:
|
||||||
|
from langchain_core.messages import HumanMessage, SystemMessage # noqa: PLC0415
|
||||||
|
messages = [
|
||||||
|
SystemMessage(content=system_text),
|
||||||
|
HumanMessage(content="Extract memory candidates as JSON."),
|
||||||
|
]
|
||||||
|
|
||||||
|
if lf:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="memory-extraction",
|
||||||
|
model=model_for_agent("memory-extractor"),
|
||||||
|
prompt=prompt_obj,
|
||||||
|
input=messages,
|
||||||
|
) as gen:
|
||||||
|
response = await llm_json.ainvoke(messages)
|
||||||
|
gen.update(output=response.content, usage=extract_usage(response))
|
||||||
|
else:
|
||||||
|
response = await llm_json.ainvoke(messages)
|
||||||
|
|
||||||
|
raw = json.loads(response.content)
|
||||||
|
result = ExtractionResult.model_validate(raw)
|
||||||
|
logger.info("memory_extraction: extracted %d candidates", len(result.candidates))
|
||||||
|
return result
|
||||||
|
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_extraction: extract_candidates failed: %s", exc)
|
||||||
|
return ExtractionResult(candidates=[])
|
||||||
|
|
||||||
|
|
||||||
|
# ── Task 2.2 — Decide action ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def decide_action(
|
||||||
|
candidate: MemoryCandidate,
|
||||||
|
existing: list[str],
|
||||||
|
) -> Literal["ADD", "UPDATE", "DELETE", "NOOP"]:
|
||||||
|
"""Decide what to do with a candidate given existing memories in the same tier.
|
||||||
|
|
||||||
|
Short-circuits to ADD without an LLM call when existing is empty (cost saving).
|
||||||
|
Never raises.
|
||||||
|
"""
|
||||||
|
if not existing:
|
||||||
|
return "ADD"
|
||||||
|
|
||||||
|
candidate_str = f"[{candidate.type}] {candidate.content}"
|
||||||
|
existing_str = "\n".join(f"- {m}" for m in existing)
|
||||||
|
|
||||||
|
template, prompt_obj = get_prompt_or_fallback("memory_decide_action", _DECIDE_FALLBACK)
|
||||||
|
|
||||||
|
if prompt_obj is not None:
|
||||||
|
try:
|
||||||
|
system_text = prompt_obj.compile(
|
||||||
|
candidate=candidate_str,
|
||||||
|
existing_memories=existing_str,
|
||||||
|
)
|
||||||
|
if isinstance(system_text, list):
|
||||||
|
system_text = "\n".join(m.get("content", "") for m in system_text if isinstance(m, dict))
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_extraction: decide compile failed: %s", exc)
|
||||||
|
system_text = template.format(candidate=candidate_str, existing_memories=existing_str)
|
||||||
|
else:
|
||||||
|
system_text = template.format(candidate=candidate_str, existing_memories=existing_str)
|
||||||
|
|
||||||
|
llm = get_agent_llm("memory-extractor", temperature=0)
|
||||||
|
lf = get_langfuse()
|
||||||
|
|
||||||
|
try:
|
||||||
|
from langchain_core.messages import HumanMessage, SystemMessage # noqa: PLC0415
|
||||||
|
messages = [
|
||||||
|
SystemMessage(content=system_text),
|
||||||
|
HumanMessage(content="Decide action."),
|
||||||
|
]
|
||||||
|
|
||||||
|
if lf:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="memory-decide-action",
|
||||||
|
model=model_for_agent("memory-extractor"),
|
||||||
|
prompt=prompt_obj,
|
||||||
|
input=messages,
|
||||||
|
) as gen:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
gen.update(output=response.content, usage=extract_usage(response))
|
||||||
|
else:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
|
||||||
|
verb = response.content.strip().upper()
|
||||||
|
if verb in ("ADD", "UPDATE", "DELETE", "NOOP"):
|
||||||
|
return verb # type: ignore[return-value]
|
||||||
|
logger.warning("memory_extraction: unexpected decide verb=%r, defaulting ADD", verb)
|
||||||
|
return "ADD"
|
||||||
|
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_extraction: decide_action failed: %s", exc)
|
||||||
|
return "ADD"
|
||||||
|
|
||||||
|
|
||||||
|
# ── Task 2.3 — Pipeline orchestrator ──────────────────────────────────────────
|
||||||
|
|
||||||
|
async def run_extraction(
|
||||||
|
db: AsyncSession,
|
||||||
|
user_id: str,
|
||||||
|
last_user_msg: str,
|
||||||
|
last_assistant_msg: str,
|
||||||
|
session_id: str | None,
|
||||||
|
) -> None:
|
||||||
|
"""Full Mem0-style extract/update pipeline for one conversation turn.
|
||||||
|
|
||||||
|
Steps:
|
||||||
|
1. Load core memory + last 5 episodes.
|
||||||
|
2. extract_candidates() → up to 5 MemoryCandidate objects.
|
||||||
|
3. For each candidate: find top-3 neighbours → decide_action() → apply.
|
||||||
|
4. Trace via Langfuse.
|
||||||
|
|
||||||
|
Never raises — wraps everything in try/except.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await _run_extraction_inner(db, user_id, last_user_msg, last_assistant_msg, session_id)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_extraction: run_extraction failed user=%s: %s", user_id, exc)
|
||||||
|
|
||||||
|
|
||||||
|
async def _run_extraction_inner(
|
||||||
|
db: AsyncSession,
|
||||||
|
user_id: str,
|
||||||
|
last_user_msg: str,
|
||||||
|
last_assistant_msg: str,
|
||||||
|
session_id: str | None,
|
||||||
|
) -> None:
|
||||||
|
from app.core.memory_middleware import MemoryMiddleware # noqa: PLC0415
|
||||||
|
|
||||||
|
middleware = MemoryMiddleware(db)
|
||||||
|
fernet = await middleware._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
logger.warning("memory_extraction: no fernet for user=%s, skipping", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
# 1. Load context
|
||||||
|
core: dict[str, str] = await middleware._load_core(user_id, fernet)
|
||||||
|
episodes: list[str] = await middleware._load_episodic(user_id, fernet, session_id=session_id)
|
||||||
|
|
||||||
|
last_turn = f"User: {last_user_msg}\nAssistant: {last_assistant_msg}"
|
||||||
|
|
||||||
|
lf = get_langfuse()
|
||||||
|
|
||||||
|
async def _run(trace_id: str | None) -> dict[str, Any]:
|
||||||
|
# 2. Extract candidates
|
||||||
|
result = await extract_candidates(last_turn, core, episodes)
|
||||||
|
if not result.candidates:
|
||||||
|
logger.info("memory_extraction: no candidates user=%s", user_id)
|
||||||
|
return {"candidates": 0, "applied": 0}
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"memory_extraction: processing %d candidates user=%s trace=%s",
|
||||||
|
len(result.candidates),
|
||||||
|
user_id,
|
||||||
|
trace_id or "-",
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3. Apply each candidate
|
||||||
|
applied = 0
|
||||||
|
actions: list[str] = []
|
||||||
|
for candidate in result.candidates:
|
||||||
|
try:
|
||||||
|
await _apply_candidate(middleware, db, user_id, fernet, candidate, trace_id)
|
||||||
|
applied += 1
|
||||||
|
actions.append(f"{candidate.type}:{candidate.target_tier}")
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory_extraction: apply failed candidate=%r user=%s: %s",
|
||||||
|
candidate.content[:80],
|
||||||
|
user_id,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"memory_extraction: applied %d/%d candidates user=%s",
|
||||||
|
applied,
|
||||||
|
len(result.candidates),
|
||||||
|
user_id,
|
||||||
|
)
|
||||||
|
return {"candidates": len(result.candidates), "applied": applied, "actions": actions}
|
||||||
|
|
||||||
|
with langfuse_context(user_id=user_id, session_id=session_id):
|
||||||
|
if lf:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="span",
|
||||||
|
name="memory-extraction-pipeline",
|
||||||
|
input={"last_turn_preview": last_turn[:200]},
|
||||||
|
) as span:
|
||||||
|
summary = await _run(trace_id=span.id)
|
||||||
|
span.update(output=summary)
|
||||||
|
try:
|
||||||
|
lf.flush()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
await _run(trace_id=None)
|
||||||
|
|
||||||
|
|
||||||
|
async def _apply_candidate(
|
||||||
|
middleware: Any,
|
||||||
|
db: AsyncSession,
|
||||||
|
user_id: str,
|
||||||
|
fernet: Any,
|
||||||
|
candidate: MemoryCandidate,
|
||||||
|
trace_id: str | None,
|
||||||
|
) -> None:
|
||||||
|
"""Fetch neighbours, decide action, apply to the appropriate tier."""
|
||||||
|
|
||||||
|
neighbours: list[str] = []
|
||||||
|
|
||||||
|
if candidate.target_tier == "core":
|
||||||
|
# For core tier: neighbours are existing core block values for similar keys.
|
||||||
|
blocks = await middleware.list_core_blocks(user_id)
|
||||||
|
neighbours = [b["value"] for b in blocks[:3]]
|
||||||
|
|
||||||
|
elif candidate.target_tier == "associative":
|
||||||
|
neighbours = await middleware.search_archival(user_id, candidate.content, top_k=3)
|
||||||
|
|
||||||
|
elif candidate.target_tier == "relational":
|
||||||
|
# Relation candidates handled specially — passed to upsert_relation directly.
|
||||||
|
# Neighbours: search by subject label if available.
|
||||||
|
neighbours = []
|
||||||
|
|
||||||
|
elif candidate.target_tier == "proactive":
|
||||||
|
neighbours = await middleware.search_recall(user_id, candidate.content, top_k=3)
|
||||||
|
|
||||||
|
action = await decide_action(candidate, neighbours)
|
||||||
|
logger.info(
|
||||||
|
"memory_extraction: candidate type=%s tier=%s action=%s",
|
||||||
|
candidate.type,
|
||||||
|
candidate.target_tier,
|
||||||
|
action,
|
||||||
|
)
|
||||||
|
|
||||||
|
if action == "NOOP":
|
||||||
|
return
|
||||||
|
|
||||||
|
if candidate.target_tier == "relational":
|
||||||
|
# Always upsert relations — decide_action skipped (no neighbour search).
|
||||||
|
if candidate.subject and candidate.predicate and candidate.object:
|
||||||
|
await _upsert_relation(
|
||||||
|
middleware, db, user_id, candidate, trace_id
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if action in ("ADD", "UPDATE"):
|
||||||
|
if candidate.target_tier == "core":
|
||||||
|
# Derive a short key from the content (first 40 chars, snake_cased).
|
||||||
|
key = _content_to_key(candidate.content)
|
||||||
|
await middleware.update_core(user_id, key, candidate.content, trace_id=trace_id)
|
||||||
|
|
||||||
|
elif candidate.target_tier == "associative":
|
||||||
|
await middleware.store_associative(user_id, candidate.content)
|
||||||
|
|
||||||
|
elif candidate.target_tier == "proactive":
|
||||||
|
await _store_proactive_stub(middleware, db, user_id, candidate, fernet)
|
||||||
|
|
||||||
|
elif action == "DELETE":
|
||||||
|
if candidate.target_tier == "core":
|
||||||
|
key = _content_to_key(candidate.content)
|
||||||
|
await middleware.delete_core(user_id, key)
|
||||||
|
|
||||||
|
|
||||||
|
def _content_to_key(content: str) -> str:
|
||||||
|
"""Derive a short snake_case key from a content string (first 40 chars)."""
|
||||||
|
import re # noqa: PLC0415
|
||||||
|
slug = re.sub(r"[^a-z0-9]+", "_", content[:40].lower()).strip("_")
|
||||||
|
return slug or "memory"
|
||||||
|
|
||||||
|
|
||||||
|
async def _upsert_relation(
|
||||||
|
middleware: Any,
|
||||||
|
db: AsyncSession,
|
||||||
|
user_id: str,
|
||||||
|
candidate: MemoryCandidate,
|
||||||
|
trace_id: str | None,
|
||||||
|
) -> None:
|
||||||
|
"""Upsert a relation row via MemoryMiddleware.upsert_relation (Phase 3)."""
|
||||||
|
await middleware.upsert_relation(
|
||||||
|
user_id=user_id,
|
||||||
|
subject=candidate.subject or "unknown",
|
||||||
|
subject_type="unknown",
|
||||||
|
predicate=candidate.predicate or "related_to",
|
||||||
|
object_=candidate.object or "unknown",
|
||||||
|
object_type="unknown",
|
||||||
|
confidence=candidate.confidence,
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"memory_extraction: upserted relation subject=%s predicate=%s object=%s",
|
||||||
|
candidate.subject,
|
||||||
|
candidate.predicate,
|
||||||
|
candidate.object,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _store_proactive_stub(
|
||||||
|
middleware: Any,
|
||||||
|
db: AsyncSession,
|
||||||
|
user_id: str,
|
||||||
|
candidate: MemoryCandidate,
|
||||||
|
fernet: Any,
|
||||||
|
) -> None:
|
||||||
|
"""Store a proactive pattern row directly (MemoryProactive model)."""
|
||||||
|
import uuid # noqa: PLC0415
|
||||||
|
from app.models import MemoryProactive # noqa: PLC0415
|
||||||
|
from app.core.memory_middleware import _encrypt # noqa: PLC0415
|
||||||
|
|
||||||
|
encrypted = _encrypt(fernet, candidate.content)
|
||||||
|
row = MemoryProactive(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
pattern_encrypted=encrypted,
|
||||||
|
confidence=candidate.confidence,
|
||||||
|
source="inferred",
|
||||||
|
)
|
||||||
|
db.add(row)
|
||||||
|
try:
|
||||||
|
await db.commit()
|
||||||
|
logger.info("memory_extraction: stored proactive pattern user=%s", user_id)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_extraction: store proactive failed: %s", exc)
|
||||||
|
await db.rollback()
|
||||||
581
app/core/memory_maintenance.py
Normal file
581
app/core/memory_maintenance.py
Normal file
@@ -0,0 +1,581 @@
|
|||||||
|
"""Memory maintenance jobs — Phase 3/5.
|
||||||
|
|
||||||
|
Three entrypoints called by the scheduler (APScheduler) registered in app/main.py:
|
||||||
|
|
||||||
|
drain_extraction_queue(db) — Free-tier batch extraction (Phase 2/5).
|
||||||
|
mine_proactive_patterns(db, user_id) — Power+ pattern mining (Phase 5).
|
||||||
|
decay_relations(db, user_id) — confidence decay + pruning for memory_relations (Phase 3).
|
||||||
|
|
||||||
|
All are safe to call manually or from tests; they never raise.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime, timedelta, timezone
|
||||||
|
|
||||||
|
from cryptography.fernet import Fernet
|
||||||
|
from sqlalchemy import select
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.core.langfuse_client import compile_prompt, extract_usage, get_langfuse, get_prompt_or_fallback
|
||||||
|
from app.models import MemoryAssociative, MemoryEpisodic, MemoryProactive, MemoryRelation, User
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Decay parameters for relations
|
||||||
|
_DECAY_FACTOR = 0.95
|
||||||
|
_DECAY_PERIOD_DAYS = 30
|
||||||
|
_PRUNE_THRESHOLD = 0.2
|
||||||
|
|
||||||
|
# Proactive pattern decay: 10 % per 7 days since last sighting
|
||||||
|
_PROACTIVE_DECAY_FACTOR = 0.9
|
||||||
|
_PROACTIVE_DECAY_PERIOD_DAYS = 7
|
||||||
|
_PROACTIVE_PRUNE_THRESHOLD = 0.2
|
||||||
|
|
||||||
|
# Mining: require at least this many episodes to attempt pattern extraction
|
||||||
|
_MIN_EPISODES_FOR_MINING = 3
|
||||||
|
_MINING_LOOKBACK_DAYS = 30
|
||||||
|
|
||||||
|
# Audit: caps to control token cost
|
||||||
|
_AUDIT_MAX_FACTS = 50
|
||||||
|
_AUDIT_MAX_LABELS = 100
|
||||||
|
|
||||||
|
|
||||||
|
async def decay_relations(db: AsyncSession, user_id: str) -> None:
|
||||||
|
"""Apply confidence decay to all relation rows for a user.
|
||||||
|
|
||||||
|
Decay rule: confidence *= 0.95 for every 30 days since last_confirmed_at.
|
||||||
|
Rows whose confidence falls below 0.2 are deleted.
|
||||||
|
|
||||||
|
Never raises — wraps in try/except.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await _decay_relations_inner(db, user_id)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: decay_relations failed user=%s: %s", user_id, exc)
|
||||||
|
|
||||||
|
|
||||||
|
async def _decay_relations_inner(db: AsyncSession, user_id: str) -> None:
|
||||||
|
result = await db.execute(
|
||||||
|
select(MemoryRelation).where(MemoryRelation.user_id == user_id)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
now = datetime.now(timezone.utc)
|
||||||
|
deleted = 0
|
||||||
|
decayed = 0
|
||||||
|
|
||||||
|
for row in rows:
|
||||||
|
reference = row.last_confirmed_at or row.created_at
|
||||||
|
if reference is None:
|
||||||
|
continue
|
||||||
|
if reference.tzinfo is None:
|
||||||
|
reference = reference.replace(tzinfo=timezone.utc)
|
||||||
|
|
||||||
|
days_elapsed = (now - reference).days
|
||||||
|
if days_elapsed < _DECAY_PERIOD_DAYS:
|
||||||
|
continue
|
||||||
|
|
||||||
|
periods = days_elapsed // _DECAY_PERIOD_DAYS
|
||||||
|
new_confidence = row.confidence * (_DECAY_FACTOR ** periods)
|
||||||
|
|
||||||
|
if new_confidence < _PRUNE_THRESHOLD:
|
||||||
|
await db.delete(row)
|
||||||
|
deleted += 1
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: pruned relation id=%s user=%s subject=%s predicate=%s "
|
||||||
|
"confidence=%.3f (below threshold)",
|
||||||
|
row.id, user_id, row.subject_label, row.predicate, new_confidence,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
row.confidence = new_confidence
|
||||||
|
decayed += 1
|
||||||
|
|
||||||
|
try:
|
||||||
|
await db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: decay_relations user=%s decayed=%d deleted=%d",
|
||||||
|
user_id, decayed, deleted,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: decay_relations commit failed user=%s: %s", user_id, exc)
|
||||||
|
await db.rollback()
|
||||||
|
|
||||||
|
|
||||||
|
async def drain_extraction_queue(db: AsyncSession) -> None:
|
||||||
|
"""Process pending ExtractionQueue rows for Free-tier users.
|
||||||
|
|
||||||
|
Each row corresponds to a stored episode that should be fed through the
|
||||||
|
Mem0-style extraction pipeline. Rows are deleted after successful processing.
|
||||||
|
Never raises — wraps in try/except.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await _drain_extraction_queue_inner(db)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: drain_extraction_queue failed: %s", exc)
|
||||||
|
|
||||||
|
|
||||||
|
async def _drain_extraction_queue_inner(db: AsyncSession) -> None:
|
||||||
|
from app.models import ExtractionQueue # noqa: PLC0415
|
||||||
|
|
||||||
|
result = await db.execute(select(ExtractionQueue))
|
||||||
|
rows = result.scalars().all()
|
||||||
|
|
||||||
|
if not rows:
|
||||||
|
logger.debug("memory_maintenance: drain_extraction_queue nothing to drain")
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info("memory_maintenance: drain_extraction_queue pending=%d", len(rows))
|
||||||
|
|
||||||
|
from app.core.memory_extraction import run_extraction # noqa: PLC0415
|
||||||
|
|
||||||
|
processed = 0
|
||||||
|
for row in rows:
|
||||||
|
try:
|
||||||
|
await run_extraction(
|
||||||
|
db=db,
|
||||||
|
user_id=row.user_id,
|
||||||
|
last_user_msg="",
|
||||||
|
last_assistant_msg="",
|
||||||
|
session_id=None,
|
||||||
|
)
|
||||||
|
await db.delete(row)
|
||||||
|
await db.commit()
|
||||||
|
processed += 1
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory_maintenance: drain failed row=%s user=%s: %s",
|
||||||
|
row.id, row.user_id, exc,
|
||||||
|
)
|
||||||
|
await db.rollback()
|
||||||
|
|
||||||
|
logger.info("memory_maintenance: drain_extraction_queue processed=%d/%d", processed, len(rows))
|
||||||
|
|
||||||
|
|
||||||
|
async def mine_proactive_patterns(db: AsyncSession, user_id: str) -> None:
|
||||||
|
"""Mine recurring behavioral patterns from last 30 days of episodes (Power+ only).
|
||||||
|
|
||||||
|
Steps:
|
||||||
|
1. Gate on proactive_mining tier feature.
|
||||||
|
2. Load + decrypt last 30 days of episodic summaries.
|
||||||
|
3. Call gpt-4o-mini to identify recurring patterns.
|
||||||
|
4. Encrypt and store each pattern in memory_proactive.
|
||||||
|
5. Apply decay to existing proactive rows.
|
||||||
|
|
||||||
|
Never raises — wraps in try/except.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await _mine_proactive_patterns_inner(db, user_id)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: mine_proactive_patterns failed user=%s: %s", user_id, exc)
|
||||||
|
|
||||||
|
|
||||||
|
async def _mine_proactive_patterns_inner(db: AsyncSession, user_id: str) -> None:
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
|
||||||
|
tier = await tier_manager.get_tier(user_id, db)
|
||||||
|
if not tier_manager.check_feature(tier, "proactive_mining"):
|
||||||
|
logger.debug("memory_maintenance: mine_proactive_patterns skipped (tier=%s)", tier)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Load user Fernet key
|
||||||
|
result = await db.execute(select(User).where(User.id == user_id))
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
if user is None or not user.encryption_key:
|
||||||
|
logger.warning("memory_maintenance: mine_proactive_patterns no encryption_key user=%s", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
fernet = Fernet(user.encryption_key.encode())
|
||||||
|
cutoff = datetime.now(timezone.utc) - timedelta(days=_MINING_LOOKBACK_DAYS)
|
||||||
|
|
||||||
|
episodes_result = await db.execute(
|
||||||
|
select(MemoryEpisodic)
|
||||||
|
.where(
|
||||||
|
MemoryEpisodic.user_id == user_id,
|
||||||
|
MemoryEpisodic.created_at >= cutoff,
|
||||||
|
)
|
||||||
|
.order_by(MemoryEpisodic.created_at.asc())
|
||||||
|
)
|
||||||
|
episode_rows = episodes_result.scalars().all()
|
||||||
|
|
||||||
|
if len(episode_rows) < _MIN_EPISODES_FOR_MINING:
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: mine_proactive_patterns skipped user=%s episodes=%d (< %d)",
|
||||||
|
user_id, len(episode_rows), _MIN_EPISODES_FOR_MINING,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
summaries: list[str] = []
|
||||||
|
for ep in episode_rows:
|
||||||
|
try:
|
||||||
|
plaintext = fernet.decrypt(ep.summary_encrypted.encode()).decode()
|
||||||
|
summaries.append(plaintext)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if not summaries:
|
||||||
|
return
|
||||||
|
|
||||||
|
patterns = await _extract_proactive_patterns(summaries)
|
||||||
|
if not patterns:
|
||||||
|
logger.info("memory_maintenance: mine_proactive_patterns user=%s no patterns extracted", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
stored = 0
|
||||||
|
for pattern_text in patterns:
|
||||||
|
try:
|
||||||
|
encrypted = fernet.encrypt(pattern_text.encode()).decode()
|
||||||
|
row = MemoryProactive(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
pattern_encrypted=encrypted,
|
||||||
|
confidence=0.7,
|
||||||
|
source="inferred",
|
||||||
|
)
|
||||||
|
db.add(row)
|
||||||
|
stored += 1
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: failed to store pattern user=%s: %s", user_id, exc)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: mine_proactive_patterns user=%s stored=%d",
|
||||||
|
user_id, stored,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: mine_proactive_patterns commit failed user=%s: %s", user_id, exc)
|
||||||
|
await db.rollback()
|
||||||
|
return
|
||||||
|
|
||||||
|
await _decay_proactive_patterns(db, user_id, fernet)
|
||||||
|
|
||||||
|
|
||||||
|
async def _extract_proactive_patterns(summaries: list[str]) -> list[str]:
|
||||||
|
"""Call memory-miner LLM to identify recurring behavioral/temporal patterns."""
|
||||||
|
from app.core.llm import get_agent_llm # noqa: PLC0415
|
||||||
|
|
||||||
|
llm = get_agent_llm("memory-miner", temperature=0)
|
||||||
|
combined = "\n---\n".join(summaries[-20:]) # cap at last 20 to control token usage
|
||||||
|
prompt = (
|
||||||
|
"You are analyzing conversation history for a personal AI secretary. "
|
||||||
|
"Identify 3-5 recurring temporal or behavioral patterns (e.g. 'always works late on Thursdays', "
|
||||||
|
"'prefers bullet-point summaries', 'frequently asks about Project Acme status'). "
|
||||||
|
"Return each pattern as a plain, short English sentence on its own line. "
|
||||||
|
"No numbering, no bullet points, no extra text.\n\n"
|
||||||
|
f"Conversation history:\n{combined}"
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
response = await llm.ainvoke(prompt)
|
||||||
|
text = response.content if hasattr(response, "content") else str(response)
|
||||||
|
lines = [line.strip() for line in str(text).splitlines() if line.strip()]
|
||||||
|
return lines[:5]
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: _extract_proactive_patterns LLM failed: %s", exc)
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
async def _decay_proactive_patterns(db: AsyncSession, user_id: str, fernet: Fernet) -> None:
|
||||||
|
"""Decay confidence of existing proactive patterns; prune below threshold."""
|
||||||
|
result = await db.execute(
|
||||||
|
select(MemoryProactive).where(MemoryProactive.user_id == user_id)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
now = datetime.now(timezone.utc)
|
||||||
|
deleted = 0
|
||||||
|
decayed = 0
|
||||||
|
|
||||||
|
for row in rows:
|
||||||
|
reference = row.created_at
|
||||||
|
if reference is None:
|
||||||
|
continue
|
||||||
|
if reference.tzinfo is None:
|
||||||
|
reference = reference.replace(tzinfo=timezone.utc)
|
||||||
|
|
||||||
|
days_elapsed = (now - reference).days
|
||||||
|
if days_elapsed < _PROACTIVE_DECAY_PERIOD_DAYS:
|
||||||
|
continue
|
||||||
|
|
||||||
|
periods = days_elapsed // _PROACTIVE_DECAY_PERIOD_DAYS
|
||||||
|
new_confidence = row.confidence * (_PROACTIVE_DECAY_FACTOR ** periods)
|
||||||
|
|
||||||
|
if new_confidence < _PROACTIVE_PRUNE_THRESHOLD:
|
||||||
|
await db.delete(row)
|
||||||
|
deleted += 1
|
||||||
|
else:
|
||||||
|
row.confidence = new_confidence
|
||||||
|
decayed += 1
|
||||||
|
|
||||||
|
try:
|
||||||
|
await db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: decay_proactive user=%s decayed=%d deleted=%d",
|
||||||
|
user_id, decayed, deleted,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: decay_proactive commit failed user=%s: %s", user_id, exc)
|
||||||
|
await db.rollback()
|
||||||
|
|
||||||
|
|
||||||
|
# ── Phase 7: weekly memory audit ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
_AUDIT_CONTRADICTIONS_FALLBACK = (
|
||||||
|
"You are auditing a personal AI assistant's memory bank. "
|
||||||
|
"Each fact has an ID in brackets. "
|
||||||
|
"Find pairs that directly contradict each other "
|
||||||
|
"(e.g. 'prefers morning meetings' vs 'never schedules before noon'). "
|
||||||
|
"For each contradiction, pick the ID to DELETE (the older or less specific one). "
|
||||||
|
'Return ONLY a valid JSON array, no markdown fences: '
|
||||||
|
'[{{"delete": "<id>", "reason": "<one line>"}}]. '
|
||||||
|
"If no contradictions, return [].\n\n"
|
||||||
|
"Facts:\n{facts}"
|
||||||
|
)
|
||||||
|
|
||||||
|
_AUDIT_CANONICALIZE_FALLBACK = (
|
||||||
|
"You are auditing entity labels in a personal AI assistant's relational memory. "
|
||||||
|
"These are names of people, companies, projects, or topics. "
|
||||||
|
"Group labels that clearly refer to the same real-world entity "
|
||||||
|
"(e.g. 'giulia', 'Giulia', 'Giulia R.' → canonical 'Giulia'). "
|
||||||
|
"Return ONLY a valid JSON array, no markdown fences: "
|
||||||
|
'[{{"canonical": "<best label>", "variants": ["<v1>", "<v2>"]}}]. '
|
||||||
|
"Only include groups with at least one variant. Singletons: omit.\n\n"
|
||||||
|
"Labels:\n{labels}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def audit_memory(db: AsyncSession, user_id: str) -> None:
|
||||||
|
"""Weekly audit: contradiction scan on associative facts + label canonicalization on relations.
|
||||||
|
|
||||||
|
Steps:
|
||||||
|
1. Decrypt up to _AUDIT_MAX_FACTS associative rows; send list to memory-auditor LLM.
|
||||||
|
2. LLM flags rows to delete (direct contradictions); hard-delete them.
|
||||||
|
3. Collect unique subject/object labels from memory_relations; ask LLM to group duplicates.
|
||||||
|
4. Rewrite variant labels to their canonical form in-place.
|
||||||
|
|
||||||
|
Never raises — wraps in try/except.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await _audit_memory_inner(db, user_id)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("memory_maintenance: audit_memory failed user=%s: %s", user_id, exc)
|
||||||
|
|
||||||
|
|
||||||
|
async def _audit_memory_inner(db: AsyncSession, user_id: str) -> None:
|
||||||
|
result = await db.execute(select(User).where(User.id == user_id))
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
if user is None or not user.encryption_key:
|
||||||
|
logger.warning("memory_maintenance: audit_memory no encryption_key user=%s", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
fernet = Fernet(user.encryption_key.encode())
|
||||||
|
await _scan_associative_contradictions(db, user_id, fernet)
|
||||||
|
await _canonicalize_relation_labels(db, user_id)
|
||||||
|
|
||||||
|
|
||||||
|
async def _scan_associative_contradictions(
|
||||||
|
db: AsyncSession,
|
||||||
|
user_id: str,
|
||||||
|
fernet: Fernet,
|
||||||
|
) -> None:
|
||||||
|
"""Decrypt associative facts, ask LLM to flag contradictions, delete superseded rows."""
|
||||||
|
result = await db.execute(
|
||||||
|
select(MemoryAssociative)
|
||||||
|
.where(MemoryAssociative.user_id == user_id)
|
||||||
|
.order_by(MemoryAssociative.updated_at.desc())
|
||||||
|
.limit(_AUDIT_MAX_FACTS)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
if len(rows) < 2:
|
||||||
|
return
|
||||||
|
|
||||||
|
id_to_text: dict[str, str] = {}
|
||||||
|
for row in rows:
|
||||||
|
try:
|
||||||
|
plaintext = fernet.decrypt(row.content_encrypted.encode()).decode()
|
||||||
|
id_to_text[row.id] = plaintext
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if len(id_to_text) < 2:
|
||||||
|
return
|
||||||
|
|
||||||
|
id_list = list(id_to_text.keys())
|
||||||
|
numbered = "\n".join(
|
||||||
|
f"{i + 1}. [{rid}] {id_to_text[rid]}" for i, rid in enumerate(id_list)
|
||||||
|
)
|
||||||
|
|
||||||
|
template, prompt_obj = get_prompt_or_fallback(
|
||||||
|
"memory_audit_contradictions", _AUDIT_CONTRADICTIONS_FALLBACK
|
||||||
|
)
|
||||||
|
system_text = compile_prompt(template, prompt_obj, facts=numbered)
|
||||||
|
|
||||||
|
from app.core.llm import get_agent_llm, model_for_agent # noqa: PLC0415
|
||||||
|
from langchain_core.messages import HumanMessage, SystemMessage # noqa: PLC0415
|
||||||
|
|
||||||
|
llm = get_agent_llm("memory-auditor", temperature=0)
|
||||||
|
lf = get_langfuse()
|
||||||
|
messages = [
|
||||||
|
SystemMessage(content=system_text),
|
||||||
|
HumanMessage(content="Audit facts for contradictions."),
|
||||||
|
]
|
||||||
|
try:
|
||||||
|
if lf:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="memory-audit-contradictions",
|
||||||
|
model=model_for_agent("memory-auditor"),
|
||||||
|
prompt=prompt_obj,
|
||||||
|
input=messages,
|
||||||
|
) as gen:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
gen.update(output=response.content, usage=extract_usage(response))
|
||||||
|
else:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
|
||||||
|
text = response.content if hasattr(response, "content") else str(response)
|
||||||
|
deletions = json.loads(text.strip())
|
||||||
|
if not isinstance(deletions, list):
|
||||||
|
return
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory_maintenance: _scan_associative_contradictions LLM/parse failed user=%s: %s",
|
||||||
|
user_id, exc,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
deleted = 0
|
||||||
|
for item in deletions:
|
||||||
|
if not isinstance(item, dict):
|
||||||
|
continue
|
||||||
|
rid = item.get("delete")
|
||||||
|
if not rid or rid not in id_to_text:
|
||||||
|
continue
|
||||||
|
result2 = await db.execute(
|
||||||
|
select(MemoryAssociative).where(
|
||||||
|
MemoryAssociative.id == rid,
|
||||||
|
MemoryAssociative.user_id == user_id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
target = result2.scalar_one_or_none()
|
||||||
|
if target:
|
||||||
|
await db.delete(target)
|
||||||
|
deleted += 1
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: audit deleted contradiction id=%s user=%s reason=%s",
|
||||||
|
rid, user_id, item.get("reason", ""),
|
||||||
|
)
|
||||||
|
|
||||||
|
if deleted:
|
||||||
|
try:
|
||||||
|
await db.commit()
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory_maintenance: audit contradiction commit failed user=%s: %s", user_id, exc
|
||||||
|
)
|
||||||
|
await db.rollback()
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: _scan_associative_contradictions user=%s deleted=%d", user_id, deleted
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _canonicalize_relation_labels(db: AsyncSession, user_id: str) -> None:
|
||||||
|
"""Group near-duplicate entity labels in memory_relations and unify to canonical form."""
|
||||||
|
result = await db.execute(
|
||||||
|
select(MemoryRelation).where(MemoryRelation.user_id == user_id)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
if not rows:
|
||||||
|
return
|
||||||
|
|
||||||
|
all_labels: set[str] = set()
|
||||||
|
for row in rows:
|
||||||
|
all_labels.add(row.subject_label)
|
||||||
|
all_labels.add(row.object_label)
|
||||||
|
|
||||||
|
labels_list = sorted(all_labels)[:_AUDIT_MAX_LABELS]
|
||||||
|
if len(labels_list) < 2:
|
||||||
|
return
|
||||||
|
|
||||||
|
labels_block = "\n".join(f"- {lbl}" for lbl in labels_list)
|
||||||
|
template, prompt_obj = get_prompt_or_fallback(
|
||||||
|
"memory_audit_canonicalize", _AUDIT_CANONICALIZE_FALLBACK
|
||||||
|
)
|
||||||
|
system_text = compile_prompt(template, prompt_obj, labels=labels_block)
|
||||||
|
|
||||||
|
from app.core.llm import get_agent_llm, model_for_agent # noqa: PLC0415
|
||||||
|
from langchain_core.messages import HumanMessage, SystemMessage # noqa: PLC0415
|
||||||
|
|
||||||
|
llm = get_agent_llm("memory-auditor", temperature=0)
|
||||||
|
lf = get_langfuse()
|
||||||
|
messages = [
|
||||||
|
SystemMessage(content=system_text),
|
||||||
|
HumanMessage(content="Canonicalize entity labels."),
|
||||||
|
]
|
||||||
|
try:
|
||||||
|
if lf:
|
||||||
|
with lf.start_as_current_observation(
|
||||||
|
as_type="generation",
|
||||||
|
name="memory-audit-canonicalize",
|
||||||
|
model=model_for_agent("memory-auditor"),
|
||||||
|
prompt=prompt_obj,
|
||||||
|
input=messages,
|
||||||
|
) as gen:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
gen.update(output=response.content, usage=extract_usage(response))
|
||||||
|
else:
|
||||||
|
response = await llm.ainvoke(messages)
|
||||||
|
|
||||||
|
text = response.content if hasattr(response, "content") else str(response)
|
||||||
|
groups = json.loads(text.strip())
|
||||||
|
if not isinstance(groups, list):
|
||||||
|
return
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory_maintenance: _canonicalize_relation_labels LLM/parse failed user=%s: %s",
|
||||||
|
user_id, exc,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Build variant → canonical map
|
||||||
|
remap: dict[str, str] = {}
|
||||||
|
for group in groups:
|
||||||
|
if not isinstance(group, dict):
|
||||||
|
continue
|
||||||
|
canonical = group.get("canonical", "")
|
||||||
|
variants = group.get("variants") or []
|
||||||
|
if not canonical:
|
||||||
|
continue
|
||||||
|
for v in variants:
|
||||||
|
if isinstance(v, str) and v != canonical:
|
||||||
|
remap[v] = canonical
|
||||||
|
|
||||||
|
if not remap:
|
||||||
|
return
|
||||||
|
|
||||||
|
updated = 0
|
||||||
|
for row in rows:
|
||||||
|
changed = False
|
||||||
|
if row.subject_label in remap:
|
||||||
|
row.subject_label = remap[row.subject_label]
|
||||||
|
changed = True
|
||||||
|
if row.object_label in remap:
|
||||||
|
row.object_label = remap[row.object_label]
|
||||||
|
changed = True
|
||||||
|
if changed:
|
||||||
|
updated += 1
|
||||||
|
|
||||||
|
if updated:
|
||||||
|
try:
|
||||||
|
await db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory_maintenance: _canonicalize_relation_labels user=%s updated=%d",
|
||||||
|
user_id, updated,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory_maintenance: canonicalize commit failed user=%s: %s", user_id, exc
|
||||||
|
)
|
||||||
|
await db.rollback()
|
||||||
733
app/core/memory_middleware.py
Normal file
733
app/core/memory_middleware.py
Normal file
@@ -0,0 +1,733 @@
|
|||||||
|
"""Memory Middleware — enrich requests with memory context and store interactions.
|
||||||
|
|
||||||
|
Four-tier memory model (MemGPT-style):
|
||||||
|
core — persistent key/value user preferences, always injected
|
||||||
|
associative — semantic similarity search via pgvector (top-k)
|
||||||
|
episodic — recent session summaries (last N)
|
||||||
|
proactive — behavioral patterns above confidence threshold
|
||||||
|
|
||||||
|
All memory content is encrypted at rest using the per-user Fernet key
|
||||||
|
stored in User.encryption_key. Decryption happens in-memory only.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
memory = MemoryMiddleware(db_session)
|
||||||
|
context = await memory.enrich_context(user_id, message)
|
||||||
|
# ... run agent ...
|
||||||
|
await memory.store_episode(user_id, session_id, message, response)
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from cryptography.fernet import Fernet, InvalidToken
|
||||||
|
from sqlalchemy import select
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.models import (
|
||||||
|
ExtractionQueue,
|
||||||
|
MemoryAssociative,
|
||||||
|
MemoryCore,
|
||||||
|
MemoryEpisodic,
|
||||||
|
MemoryProactive,
|
||||||
|
MemoryRelation,
|
||||||
|
User,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _now() -> datetime:
|
||||||
|
return datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
|
||||||
|
# Tuning constants
|
||||||
|
_ASSOCIATIVE_TOP_K = 5
|
||||||
|
_EPISODIC_RECENT_N = 10
|
||||||
|
_PROACTIVE_CONFIDENCE_THRESHOLD = 0.6
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryMiddleware:
|
||||||
|
"""Enrich orchestrator context with memory and persist interactions after."""
|
||||||
|
|
||||||
|
def __init__(self, db: AsyncSession) -> None:
|
||||||
|
self._db = db
|
||||||
|
|
||||||
|
# ── Public API ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def enrich_context(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
message: str,
|
||||||
|
trace_id: str | None = None,
|
||||||
|
session_id: str | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Build memory context dict to inject into the orchestrator before LLM call.
|
||||||
|
|
||||||
|
Returns a dict with keys:
|
||||||
|
core_memory — {key: plaintext_value, ...}
|
||||||
|
associative_memory — [plaintext_content, ...] (top-k by keyword match)
|
||||||
|
episodic_memory — [plaintext_summary, ...] (most recent N)
|
||||||
|
proactive_hints — [plaintext_pattern, ...] (above threshold)
|
||||||
|
relational_memory — ["subject --predicate--> object", ...] (top 10, Pro+)
|
||||||
|
"""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
user_dbg = await self._get_user_debug(user_id)
|
||||||
|
user_tier: str = user_dbg.get("tier") or "free"
|
||||||
|
|
||||||
|
core = await self._load_core(user_id, fernet)
|
||||||
|
associative = await self._load_associative(user_id, message, fernet, user_tier=user_tier)
|
||||||
|
episodic = await self._load_episodic(user_id, fernet, session_id=session_id)
|
||||||
|
proactive = await self._load_proactive(user_id, fernet)
|
||||||
|
relational = await self._load_relational(user_id, user_tier=user_tier)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"memory: enrich_context trace=%s user=%s tier=%s core=%d associative=%d episodic=%d proactive=%d relational=%d",
|
||||||
|
trace_id or "-",
|
||||||
|
user_id,
|
||||||
|
user_tier,
|
||||||
|
len(core),
|
||||||
|
len(associative),
|
||||||
|
len(episodic),
|
||||||
|
len(proactive),
|
||||||
|
len(relational),
|
||||||
|
)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"core_memory": core,
|
||||||
|
"associative_memory": associative,
|
||||||
|
"episodic_memory": episodic,
|
||||||
|
"proactive_hints": proactive,
|
||||||
|
"relational_memory": relational,
|
||||||
|
}
|
||||||
|
|
||||||
|
async def store_episode(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
session_id: str,
|
||||||
|
message: str,
|
||||||
|
response: str,
|
||||||
|
trace_id: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Summarise and store a completed interaction in episodic memory.
|
||||||
|
|
||||||
|
The summary is a simple heuristic concatenation (no LLM call) to keep
|
||||||
|
latency low. After committing the episode row, dispatches the Mem0-style
|
||||||
|
extraction pipeline:
|
||||||
|
- Pro/Power/Team → asyncio.create_task (fire-and-forget, realtime).
|
||||||
|
- Free → enqueue an ExtractionQueue row for the daily cron.
|
||||||
|
"""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
summary = f"User: {message[:200]}\nAssistant: {response[:200]}"
|
||||||
|
encrypted = _encrypt(fernet, summary)
|
||||||
|
|
||||||
|
episode = MemoryEpisodic(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
summary_encrypted=encrypted,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
self._db.add(episode)
|
||||||
|
episode_id: str = episode.id
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
user_dbg = await self._get_user_debug(user_id)
|
||||||
|
tier = user_dbg.get("tier") or "free"
|
||||||
|
logger.info(
|
||||||
|
"memory: store_episode trace=%s user=%s tier=%s session=%s",
|
||||||
|
trace_id or "-",
|
||||||
|
user_id,
|
||||||
|
tier,
|
||||||
|
session_id,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("memory: store_episode failed user=%s: %s", user_id, exc)
|
||||||
|
await self._db.rollback()
|
||||||
|
return
|
||||||
|
|
||||||
|
# ── Dispatch extraction pipeline (Phase 2) ────────────────────────────
|
||||||
|
await self._dispatch_extraction(
|
||||||
|
user_id=user_id,
|
||||||
|
episode_id=episode_id,
|
||||||
|
last_user_msg=message,
|
||||||
|
last_assistant_msg=response,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _dispatch_extraction(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
episode_id: str,
|
||||||
|
last_user_msg: str,
|
||||||
|
last_assistant_msg: str,
|
||||||
|
session_id: str | None,
|
||||||
|
) -> None:
|
||||||
|
"""Route extraction to realtime task or batch queue based on user tier."""
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
|
||||||
|
tier = await tier_manager.get_tier(user_id, self._db)
|
||||||
|
|
||||||
|
if tier_manager.check_feature(tier, "realtime_extraction"):
|
||||||
|
# Pro/Power/Team: fire-and-forget in the background.
|
||||||
|
# Must open a fresh session — request session closes after handler returns.
|
||||||
|
from app.core.memory_extraction import run_extraction # noqa: PLC0415
|
||||||
|
from app.db import async_session # noqa: PLC0415
|
||||||
|
|
||||||
|
async def _task() -> None:
|
||||||
|
try:
|
||||||
|
async with async_session() as fresh_db:
|
||||||
|
await run_extraction(
|
||||||
|
db=fresh_db,
|
||||||
|
user_id=user_id,
|
||||||
|
last_user_msg=last_user_msg,
|
||||||
|
last_assistant_msg=last_assistant_msg,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory: extraction task failed user=%s: %s", user_id, exc
|
||||||
|
)
|
||||||
|
|
||||||
|
asyncio.create_task(_task())
|
||||||
|
logger.info("memory: realtime extraction dispatched user=%s", user_id)
|
||||||
|
else:
|
||||||
|
# Free tier: enqueue for daily batch cron.
|
||||||
|
queue_row = ExtractionQueue(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
episode_id=episode_id,
|
||||||
|
)
|
||||||
|
self._db.add(queue_row)
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory: extraction enqueued (batch) user=%s episode=%s",
|
||||||
|
user_id,
|
||||||
|
episode_id,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory: extraction queue insert failed user=%s: %s", user_id, exc
|
||||||
|
)
|
||||||
|
await self._db.rollback()
|
||||||
|
|
||||||
|
async def update_core(self, user_id: str, key: str, value: str, trace_id: str | None = None) -> None:
|
||||||
|
"""Upsert a core memory key/value for a user."""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
encrypted = _encrypt(fernet, value)
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryCore).where(
|
||||||
|
MemoryCore.user_id == user_id,
|
||||||
|
MemoryCore.key == key,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
existing = result.scalar_one_or_none()
|
||||||
|
if existing is not None:
|
||||||
|
existing.value_encrypted = encrypted
|
||||||
|
else:
|
||||||
|
self._db.add(MemoryCore(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
key=key,
|
||||||
|
value_encrypted=encrypted,
|
||||||
|
))
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
user_dbg = await self._get_user_debug(user_id)
|
||||||
|
logger.info(
|
||||||
|
"memory: update_core trace=%s user=%s tier=%s key=%s",
|
||||||
|
trace_id or "-",
|
||||||
|
user_id,
|
||||||
|
user_dbg.get("tier") or "-",
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("memory: update_core failed user=%s key=%s: %s", user_id, key, exc)
|
||||||
|
await self._db.rollback()
|
||||||
|
|
||||||
|
async def list_core_blocks(self, user_id: str) -> list[dict[str, str]]:
|
||||||
|
"""Return core memory as editable blocks (label/value)."""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return []
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryCore)
|
||||||
|
.where(MemoryCore.user_id == user_id)
|
||||||
|
.order_by(MemoryCore.key.asc())
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out: list[dict[str, str]] = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.value_encrypted)
|
||||||
|
if plaintext is not None:
|
||||||
|
out.append({"label": row.key, "value": plaintext})
|
||||||
|
logger.debug("memory: list_core_blocks user=%s count=%d", user_id, len(out))
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def get_core_block(self, user_id: str, label: str) -> str | None:
|
||||||
|
"""Return a single core memory block value by label."""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryCore).where(
|
||||||
|
MemoryCore.user_id == user_id,
|
||||||
|
MemoryCore.key == label,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
row = result.scalar_one_or_none()
|
||||||
|
if row is None:
|
||||||
|
logger.debug("memory: get_core_block user=%s label=%s found=0", user_id, label)
|
||||||
|
return None
|
||||||
|
value = _safe_decrypt(fernet, row.value_encrypted)
|
||||||
|
logger.debug("memory: get_core_block user=%s label=%s found=%d", user_id, label, 1 if value is not None else 0)
|
||||||
|
return value
|
||||||
|
|
||||||
|
async def delete_core(self, user_id: str, label: str) -> bool:
|
||||||
|
"""Delete a core memory block by label. Returns True if deleted."""
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryCore).where(
|
||||||
|
MemoryCore.user_id == user_id,
|
||||||
|
MemoryCore.key == label,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
row = result.scalar_one_or_none()
|
||||||
|
if row is None:
|
||||||
|
logger.debug("memory: delete_core user=%s label=%s found=0", user_id, label)
|
||||||
|
return False
|
||||||
|
|
||||||
|
await self._db.delete(row)
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
logger.info("memory: delete_core user=%s label=%s", user_id, label)
|
||||||
|
return True
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("memory: delete_core failed user=%s label=%s: %s", user_id, label, exc)
|
||||||
|
await self._db.rollback()
|
||||||
|
return False
|
||||||
|
|
||||||
|
async def append_core(self, user_id: str, label: str, content: str) -> None:
|
||||||
|
"""Append content to a core block, creating it if missing."""
|
||||||
|
current = await self.get_core_block(user_id, label)
|
||||||
|
if current is None:
|
||||||
|
await self.update_core(user_id, label, content)
|
||||||
|
logger.info("memory: append_core user=%s label=%s created=1", user_id, label)
|
||||||
|
return
|
||||||
|
await self.update_core(user_id, label, f"{current}\n{content}")
|
||||||
|
logger.info("memory: append_core user=%s label=%s created=0", user_id, label)
|
||||||
|
|
||||||
|
async def replace_core(self, user_id: str, label: str, old: str, new: str) -> bool:
|
||||||
|
"""Replace one exact string inside a core block. Returns False if not found."""
|
||||||
|
current = await self.get_core_block(user_id, label)
|
||||||
|
if current is None or old not in current:
|
||||||
|
logger.debug("memory: replace_core user=%s label=%s changed=0", user_id, label)
|
||||||
|
return False
|
||||||
|
await self.update_core(user_id, label, current.replace(old, new, 1))
|
||||||
|
logger.info("memory: replace_core user=%s label=%s changed=1", user_id, label)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def store_associative(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
content: str,
|
||||||
|
entity_type: str | None = None,
|
||||||
|
entity_id: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Store associative memory; embed if user tier has real_embeddings."""
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
from app.core.embeddings import embed_text # noqa: PLC0415
|
||||||
|
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
encrypted = _encrypt(fernet, content)
|
||||||
|
|
||||||
|
user_dbg = await self._get_user_debug(user_id)
|
||||||
|
user_tier = user_dbg.get("tier") or "free"
|
||||||
|
|
||||||
|
embedding: list[float] | None = None
|
||||||
|
if tier_manager.check_feature(user_tier, "real_embeddings"):
|
||||||
|
embedding = await embed_text(content)
|
||||||
|
|
||||||
|
row = MemoryAssociative(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
content_encrypted=encrypted,
|
||||||
|
embedding=embedding,
|
||||||
|
entity_type=entity_type,
|
||||||
|
entity_id=entity_id,
|
||||||
|
)
|
||||||
|
self._db.add(row)
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory: store_associative user=%s embedded=%s",
|
||||||
|
user_id,
|
||||||
|
embedding is not None,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("memory: store_associative failed user=%s: %s", user_id, exc)
|
||||||
|
await self._db.rollback()
|
||||||
|
|
||||||
|
async def upsert_relation(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
subject: str,
|
||||||
|
subject_type: str,
|
||||||
|
predicate: str,
|
||||||
|
object_: str,
|
||||||
|
object_type: str,
|
||||||
|
*,
|
||||||
|
confidence: float = 0.7,
|
||||||
|
source_episode_id: str | None = None,
|
||||||
|
notes: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Insert or update a relation row. Matches on (user_id, subject_label, predicate, object_label).
|
||||||
|
|
||||||
|
subject_label / object_label are plaintext entity identifiers — not encrypted.
|
||||||
|
notes is optional; encrypted with user Fernet if provided.
|
||||||
|
"""
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
|
||||||
|
user_dbg = await self._get_user_debug(user_id)
|
||||||
|
user_tier = user_dbg.get("tier") or "free"
|
||||||
|
if not tier_manager.check_feature(user_tier, "relational_memory"):
|
||||||
|
logger.debug("memory: upsert_relation skipped (tier=%s no relational_memory)", user_tier)
|
||||||
|
return
|
||||||
|
|
||||||
|
notes_encrypted: bytes | None = None
|
||||||
|
if notes:
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet:
|
||||||
|
notes_encrypted = fernet.encrypt(notes.encode())
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryRelation).where(
|
||||||
|
MemoryRelation.user_id == user_id,
|
||||||
|
MemoryRelation.subject_label == subject,
|
||||||
|
MemoryRelation.predicate == predicate,
|
||||||
|
MemoryRelation.object_label == object_,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
existing = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if existing is not None:
|
||||||
|
existing.subject_type = subject_type
|
||||||
|
existing.object_type = object_type
|
||||||
|
existing.confidence = confidence
|
||||||
|
existing.last_confirmed_at = _now()
|
||||||
|
if notes_encrypted is not None:
|
||||||
|
existing.notes_encrypted = notes_encrypted
|
||||||
|
else:
|
||||||
|
self._db.add(MemoryRelation(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
subject_label=subject,
|
||||||
|
subject_type=subject_type,
|
||||||
|
predicate=predicate,
|
||||||
|
object_label=object_,
|
||||||
|
object_type=object_type,
|
||||||
|
confidence=confidence,
|
||||||
|
source_episode_id=source_episode_id,
|
||||||
|
notes_encrypted=notes_encrypted,
|
||||||
|
))
|
||||||
|
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
logger.info(
|
||||||
|
"memory: upsert_relation user=%s subject=%s predicate=%s object=%s",
|
||||||
|
user_id, subject, predicate, object_,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("memory: upsert_relation failed user=%s: %s", user_id, exc)
|
||||||
|
await self._db.rollback()
|
||||||
|
|
||||||
|
async def query_relations(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
subject: str | None = None,
|
||||||
|
predicate: str | None = None,
|
||||||
|
object_: str | None = None,
|
||||||
|
limit: int = 20,
|
||||||
|
) -> list[MemoryRelation]:
|
||||||
|
"""Query relation rows for a user with optional filters."""
|
||||||
|
q = select(MemoryRelation).where(MemoryRelation.user_id == user_id)
|
||||||
|
if subject is not None:
|
||||||
|
q = q.where(MemoryRelation.subject_label == subject)
|
||||||
|
if predicate is not None:
|
||||||
|
q = q.where(MemoryRelation.predicate == predicate)
|
||||||
|
if object_ is not None:
|
||||||
|
q = q.where(MemoryRelation.object_label == object_)
|
||||||
|
q = q.order_by(MemoryRelation.confidence.desc()).limit(limit)
|
||||||
|
result = await self._db.execute(q)
|
||||||
|
return list(result.scalars().all())
|
||||||
|
|
||||||
|
async def insert_archival(self, user_id: str, content: str, source: str = "manual") -> None:
|
||||||
|
"""Insert a long-term archival memory entry."""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
encrypted = _encrypt(fernet, content)
|
||||||
|
row = MemoryAssociative(
|
||||||
|
id=str(uuid.uuid4()),
|
||||||
|
user_id=user_id,
|
||||||
|
content_encrypted=encrypted,
|
||||||
|
embedding=None,
|
||||||
|
entity_type=source,
|
||||||
|
entity_id=None,
|
||||||
|
)
|
||||||
|
self._db.add(row)
|
||||||
|
try:
|
||||||
|
await self._db.commit()
|
||||||
|
logger.info("memory: insert_archival user=%s source=%s", user_id, source)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("memory: insert_archival failed user=%s: %s", user_id, exc)
|
||||||
|
await self._db.rollback()
|
||||||
|
|
||||||
|
async def search_archival(self, user_id: str, query: str, top_k: int = 5) -> list[str]:
|
||||||
|
"""Search archival memory (keyword fallback; semantic ranking can replace this)."""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return []
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryAssociative)
|
||||||
|
.where(MemoryAssociative.user_id == user_id)
|
||||||
|
.order_by(MemoryAssociative.updated_at.desc())
|
||||||
|
.limit(100)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
needle = query.strip().lower()
|
||||||
|
out: list[str] = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.content_encrypted)
|
||||||
|
if plaintext is None:
|
||||||
|
continue
|
||||||
|
if not needle or needle in plaintext.lower():
|
||||||
|
out.append(plaintext)
|
||||||
|
if len(out) >= max(top_k, 1):
|
||||||
|
break
|
||||||
|
logger.info("memory: search_archival user=%s query=%s hits=%d", user_id, query[:80], len(out))
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def search_recall(self, user_id: str, query: str, top_k: int = 5) -> list[str]:
|
||||||
|
"""Search recall memory (episodic summaries) by keyword."""
|
||||||
|
fernet = await self._get_fernet(user_id)
|
||||||
|
if fernet is None:
|
||||||
|
return []
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryEpisodic)
|
||||||
|
.where(MemoryEpisodic.user_id == user_id)
|
||||||
|
.order_by(MemoryEpisodic.created_at.desc())
|
||||||
|
.limit(100)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
needle = query.strip().lower()
|
||||||
|
out: list[str] = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.summary_encrypted)
|
||||||
|
if plaintext is None:
|
||||||
|
continue
|
||||||
|
if not needle or needle in plaintext.lower():
|
||||||
|
out.append(plaintext)
|
||||||
|
if len(out) >= max(top_k, 1):
|
||||||
|
break
|
||||||
|
logger.info("memory: search_recall user=%s query=%s hits=%d", user_id, query[:80], len(out))
|
||||||
|
return out
|
||||||
|
|
||||||
|
# ── Private helpers ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def _get_fernet(self, user_id: str) -> Fernet | None:
|
||||||
|
"""Load the user's Fernet key from DB. Returns None if missing."""
|
||||||
|
result = await self._db.execute(select(User).where(User.id == user_id))
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
if user is None or not user.encryption_key:
|
||||||
|
logger.warning("memory: no encryption_key for user=%s", user_id)
|
||||||
|
return None
|
||||||
|
return Fernet(user.encryption_key.encode())
|
||||||
|
|
||||||
|
async def _get_user_debug(self, user_id: str) -> dict[str, str | None]:
|
||||||
|
"""Load lightweight user debug fields for trace logs."""
|
||||||
|
from app.config.settings import settings # noqa: PLC0415
|
||||||
|
from app.models import Subscription # noqa: PLC0415
|
||||||
|
|
||||||
|
result = await self._db.execute(select(User).where(User.id == user_id))
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
if user is None:
|
||||||
|
return {"tier": None}
|
||||||
|
|
||||||
|
sub_result = await self._db.execute(
|
||||||
|
select(Subscription.tier).where(Subscription.user_id == user_id)
|
||||||
|
)
|
||||||
|
sub_tier: str | None = sub_result.scalar_one_or_none()
|
||||||
|
if sub_tier:
|
||||||
|
tier = sub_tier
|
||||||
|
elif settings.ENV == "dev":
|
||||||
|
tier = "power"
|
||||||
|
else:
|
||||||
|
tier = user.tier or "free"
|
||||||
|
|
||||||
|
return {"tier": tier}
|
||||||
|
|
||||||
|
async def _load_core(self, user_id: str, fernet: Fernet) -> dict[str, str]:
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryCore).where(MemoryCore.user_id == user_id)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out: dict[str, str] = {}
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.value_encrypted)
|
||||||
|
if plaintext is not None:
|
||||||
|
out[row.key] = plaintext
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def _load_associative(
|
||||||
|
self, user_id: str, message: str, fernet: Fernet, *, user_tier: str = "free"
|
||||||
|
) -> list[str]:
|
||||||
|
"""Load top-k associative memories.
|
||||||
|
|
||||||
|
Pro+: pgvector cosine similarity on the message embedding (real_embeddings feature).
|
||||||
|
Free / embedding failure: keyword-ordered fallback (most recent rows).
|
||||||
|
"""
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
from app.core.embeddings import embed_text # noqa: PLC0415
|
||||||
|
|
||||||
|
if tier_manager.check_feature(user_tier, "real_embeddings"):
|
||||||
|
vec = await embed_text(message)
|
||||||
|
if vec is not None:
|
||||||
|
try:
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryAssociative)
|
||||||
|
.where(
|
||||||
|
MemoryAssociative.user_id == user_id,
|
||||||
|
MemoryAssociative.embedding.isnot(None),
|
||||||
|
)
|
||||||
|
.order_by(MemoryAssociative.embedding.cosine_distance(vec))
|
||||||
|
.limit(_ASSOCIATIVE_TOP_K)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out: list[str] = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.content_encrypted)
|
||||||
|
if plaintext is not None:
|
||||||
|
out.append(plaintext)
|
||||||
|
logger.info(
|
||||||
|
"memory: _load_associative user=%s mode=vector hits=%d",
|
||||||
|
user_id,
|
||||||
|
len(out),
|
||||||
|
)
|
||||||
|
return out
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"memory: vector search failed user=%s, falling back to keyword: %s",
|
||||||
|
user_id,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Keyword fallback: most recent rows
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryAssociative)
|
||||||
|
.where(MemoryAssociative.user_id == user_id)
|
||||||
|
.order_by(MemoryAssociative.updated_at.desc())
|
||||||
|
.limit(_ASSOCIATIVE_TOP_K)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.content_encrypted)
|
||||||
|
if plaintext is not None:
|
||||||
|
out.append(plaintext)
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def _load_episodic(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
fernet: Fernet,
|
||||||
|
session_id: str | None = None,
|
||||||
|
) -> list[str]:
|
||||||
|
query = select(MemoryEpisodic).where(MemoryEpisodic.user_id == user_id)
|
||||||
|
if session_id:
|
||||||
|
query = query.where(MemoryEpisodic.session_id == session_id)
|
||||||
|
result = await self._db.execute(
|
||||||
|
query
|
||||||
|
.order_by(MemoryEpisodic.created_at.desc())
|
||||||
|
.limit(_EPISODIC_RECENT_N)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out: list[str] = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.summary_encrypted)
|
||||||
|
if plaintext is not None:
|
||||||
|
out.append(plaintext)
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def _load_relational(self, user_id: str, *, user_tier: str = "free") -> list[str]:
|
||||||
|
"""Return top-10 relation strings for Pro+ users; empty list for Free."""
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
|
||||||
|
if not tier_manager.check_feature(user_tier, "relational_memory"):
|
||||||
|
return []
|
||||||
|
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryRelation)
|
||||||
|
.where(MemoryRelation.user_id == user_id)
|
||||||
|
.order_by(MemoryRelation.confidence.desc())
|
||||||
|
.limit(10)
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out = [
|
||||||
|
f"{r.subject_label} --{r.predicate}--> {r.object_label}"
|
||||||
|
for r in rows
|
||||||
|
]
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def _load_proactive(self, user_id: str, fernet: Fernet) -> list[str]:
|
||||||
|
result = await self._db.execute(
|
||||||
|
select(MemoryProactive)
|
||||||
|
.where(
|
||||||
|
MemoryProactive.user_id == user_id,
|
||||||
|
MemoryProactive.confidence >= _PROACTIVE_CONFIDENCE_THRESHOLD,
|
||||||
|
)
|
||||||
|
.order_by(MemoryProactive.confidence.desc())
|
||||||
|
)
|
||||||
|
rows = result.scalars().all()
|
||||||
|
out: list[str] = []
|
||||||
|
for row in rows:
|
||||||
|
plaintext = _safe_decrypt(fernet, row.pattern_encrypted)
|
||||||
|
if plaintext is not None:
|
||||||
|
out.append(plaintext)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
# ── Encryption helpers ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def _encrypt(fernet: Fernet, plaintext: str) -> str:
|
||||||
|
return fernet.encrypt(plaintext.encode()).decode()
|
||||||
|
|
||||||
|
|
||||||
|
def _safe_decrypt(fernet: Fernet, ciphertext: str) -> str | None:
|
||||||
|
"""Decrypt and return plaintext, or None on error (corrupted/wrong key)."""
|
||||||
|
try:
|
||||||
|
return fernet.decrypt(ciphertext.encode()).decode()
|
||||||
|
except (InvalidToken, Exception) as exc:
|
||||||
|
logger.warning("memory: decrypt failed: %s", exc)
|
||||||
|
return None
|
||||||
51
app/core/note_summarizer.py
Normal file
51
app/core/note_summarizer.py
Normal file
@@ -0,0 +1,51 @@
|
|||||||
|
"""Note summarizer — generates a compact AI summary for a note.
|
||||||
|
|
||||||
|
Called fire-and-forget from create_note / update_note tools so the
|
||||||
|
``notes.ai_summary`` column stays current without blocking the agent loop.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from langchain_core.messages import HumanMessage, SystemMessage
|
||||||
|
|
||||||
|
from app.core.langfuse_client import get_prompt_or_fallback
|
||||||
|
from app.core.llm import get_agent_llm
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_FALLBACK_PROMPT = """\
|
||||||
|
Summarize this note in <=250 characters. Be terse and dense.
|
||||||
|
Keep proper nouns, dates, decisions, and action items.
|
||||||
|
Do not start with "This note".
|
||||||
|
Respond with the summary text only — no intro, no labels.
|
||||||
|
|
||||||
|
Title: {title}
|
||||||
|
Content: {content}"""
|
||||||
|
|
||||||
|
_MAX_CONTENT_CHARS = 4000
|
||||||
|
|
||||||
|
|
||||||
|
async def generate_note_summary(title: str, content: str) -> str:
|
||||||
|
"""Return a <=250-char summary of *title* + *content*.
|
||||||
|
|
||||||
|
Uses the Langfuse ``note_summary`` prompt (hot-swappable) with a local
|
||||||
|
fallback. Truncates *content* to 4000 chars before sending to avoid
|
||||||
|
token waste on large notes.
|
||||||
|
"""
|
||||||
|
template, _ = get_prompt_or_fallback("note_summary", _FALLBACK_PROMPT)
|
||||||
|
trimmed = content[:_MAX_CONTENT_CHARS]
|
||||||
|
system_prompt = template.format(title=title, content=trimmed)
|
||||||
|
|
||||||
|
try:
|
||||||
|
llm = get_agent_llm("note-summarizer")
|
||||||
|
response = await llm.ainvoke([
|
||||||
|
SystemMessage(content=system_prompt),
|
||||||
|
HumanMessage(content="Generate the summary."),
|
||||||
|
])
|
||||||
|
text = response.content if isinstance(response.content, str) else ""
|
||||||
|
return text.strip()[:250]
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("note_summarizer: failed to generate summary: %s", exc)
|
||||||
|
return ""
|
||||||
@@ -1,166 +0,0 @@
|
|||||||
"""Orchestrator — LLM-based intent router and agent pipeline."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import json
|
|
||||||
from typing import Any, AsyncGenerator
|
|
||||||
|
|
||||||
from langchain_core.messages import HumanMessage, SystemMessage
|
|
||||||
|
|
||||||
from app.core.agent_registry import AgentRegistry
|
|
||||||
from app.core.llm import get_router_llm
|
|
||||||
from app.core.agent_registry import registry as _default_registry
|
|
||||||
from app.schemas import ChatRequest, ChatResponse, ExecutionPlan
|
|
||||||
|
|
||||||
_FALLBACK_AGENT = "task_agent"
|
|
||||||
|
|
||||||
_CLASSIFY_SYSTEM = (
|
|
||||||
"You are an intent classifier. Given the user message and context, decide "
|
|
||||||
"which agent to route to.\n"
|
|
||||||
"Available agents: {agents}\n"
|
|
||||||
"Respond with just the agent name, nothing else."
|
|
||||||
)
|
|
||||||
|
|
||||||
_SYNTHESIZE_HUMAN = (
|
|
||||||
"Combine the following agent results into one coherent response.\n\n"
|
|
||||||
"Agent results:\n{results}\n\n"
|
|
||||||
"Original message: {message}"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _make_llm():
|
|
||||||
return get_router_llm()
|
|
||||||
|
|
||||||
|
|
||||||
async def classify_intent(
|
|
||||||
message: str,
|
|
||||||
context: dict[str, Any],
|
|
||||||
reg: AgentRegistry,
|
|
||||||
) -> str:
|
|
||||||
"""Use gpt-4o-mini to classify intent and return the matching agent name.
|
|
||||||
|
|
||||||
Falls back to ``task_agent`` when the registry is empty or the model
|
|
||||||
returns a name that is not registered.
|
|
||||||
"""
|
|
||||||
agents = reg.list_agents()
|
|
||||||
if not agents:
|
|
||||||
return _FALLBACK_AGENT
|
|
||||||
|
|
||||||
system = _CLASSIFY_SYSTEM.format(agents=json.dumps(agents))
|
|
||||||
# Truncate context to keep the classification prompt short
|
|
||||||
human = f"Message: {message}\nContext summary: {json.dumps(context)[:500]}"
|
|
||||||
|
|
||||||
llm = _make_llm()
|
|
||||||
response = await llm.ainvoke(
|
|
||||||
[SystemMessage(content=system), HumanMessage(content=human)]
|
|
||||||
)
|
|
||||||
|
|
||||||
agent_name = str(response.content).strip().lower()
|
|
||||||
known = {a["name"] for a in agents}
|
|
||||||
return agent_name if agent_name in known else _FALLBACK_AGENT
|
|
||||||
|
|
||||||
|
|
||||||
async def route_single(
|
|
||||||
agent_name: str,
|
|
||||||
message: str,
|
|
||||||
context: dict[str, Any],
|
|
||||||
reg: AgentRegistry,
|
|
||||||
) -> ChatResponse:
|
|
||||||
"""Route to a single agent and wrap the result in a ``ChatResponse``."""
|
|
||||||
response_text = await reg.call_agent(agent_name, message, context)
|
|
||||||
return ChatResponse(response=response_text)
|
|
||||||
|
|
||||||
|
|
||||||
async def route_pipeline(
|
|
||||||
agent_names: list[str],
|
|
||||||
message: str,
|
|
||||||
context: dict[str, Any],
|
|
||||||
reg: AgentRegistry,
|
|
||||||
) -> ChatResponse:
|
|
||||||
"""Execute agents sequentially; each agent receives previous results in context.
|
|
||||||
|
|
||||||
A final LLM synthesis call merges all results into one coherent response.
|
|
||||||
"""
|
|
||||||
previous_results: list[str] = []
|
|
||||||
|
|
||||||
for agent_name in agent_names:
|
|
||||||
ctx = {**context, "previous_results": list(previous_results)}
|
|
||||||
result = await reg.call_agent(agent_name, message, ctx)
|
|
||||||
previous_results.append(result)
|
|
||||||
|
|
||||||
results_str = "\n\n".join(
|
|
||||||
f"[{name}]: {res}" for name, res in zip(agent_names, previous_results)
|
|
||||||
)
|
|
||||||
human = _SYNTHESIZE_HUMAN.format(results=results_str, message=message)
|
|
||||||
llm = _make_llm()
|
|
||||||
synthesis = await llm.ainvoke([HumanMessage(content=human)])
|
|
||||||
return ChatResponse(response=str(synthesis.content))
|
|
||||||
|
|
||||||
|
|
||||||
def _build_plan(agent_name: str, message: str) -> ExecutionPlan:
|
|
||||||
"""Build an ``ExecutionPlan`` for the resolved agent.
|
|
||||||
|
|
||||||
Uses ``ExecutionPlanBuilder`` with the server-side template registry.
|
|
||||||
If a default template exists for the agent, an LLM step is emitted;
|
|
||||||
otherwise a plain ``handle`` action step is used.
|
|
||||||
"""
|
|
||||||
from app.core.execution_plan import ExecutionPlanBuilder, template_registry
|
|
||||||
|
|
||||||
template_id = f"tpl_{agent_name}_default"
|
|
||||||
builder = ExecutionPlanBuilder(agent_name)
|
|
||||||
if template_registry.has(template_id):
|
|
||||||
builder.add_llm_step(template_id, {"message": message})
|
|
||||||
else:
|
|
||||||
builder.add_step("handle", {"message": message})
|
|
||||||
return builder.build()
|
|
||||||
|
|
||||||
|
|
||||||
async def orchestrate(
|
|
||||||
request: ChatRequest,
|
|
||||||
reg: AgentRegistry | None = None,
|
|
||||||
) -> ChatResponse | ExecutionPlan:
|
|
||||||
"""Main orchestration entry point.
|
|
||||||
|
|
||||||
* Classifies the user's intent to select an agent.
|
|
||||||
* ``execution_mode == 'direct'``: routes to the agent and returns a
|
|
||||||
``ChatResponse``.
|
|
||||||
* ``execution_mode == 'plan'``: returns an ``ExecutionPlan`` with the
|
|
||||||
resolved agent and a template-ID-only step (prompt IP stays server-side).
|
|
||||||
"""
|
|
||||||
if reg is None:
|
|
||||||
reg = _default_registry
|
|
||||||
|
|
||||||
context = request.context.model_dump()
|
|
||||||
agent_name = await classify_intent(request.message, context, reg)
|
|
||||||
|
|
||||||
if request.execution_mode == "direct":
|
|
||||||
return await route_single(agent_name, request.message, context, reg)
|
|
||||||
|
|
||||||
# plan mode — return plan, do not execute
|
|
||||||
return _build_plan(agent_name, request.message)
|
|
||||||
|
|
||||||
|
|
||||||
async def orchestrate_stream(
|
|
||||||
request: ChatRequest,
|
|
||||||
reg: AgentRegistry | None = None,
|
|
||||||
) -> AsyncGenerator[str, None]:
|
|
||||||
"""Streaming orchestration — yields plain text chunks only.
|
|
||||||
|
|
||||||
The WebSocket handler in ``app/api/routes/chat.py`` is responsible for
|
|
||||||
wrapping each chunk in a ``text_chunk`` frame and sending the final
|
|
||||||
``final`` frame once the generator is exhausted.
|
|
||||||
|
|
||||||
Agents do not yet support token-level streaming; the full response is
|
|
||||||
fetched first (which may involve multiple WS round-trips for tool calls),
|
|
||||||
then emitted in fixed-size chunks.
|
|
||||||
"""
|
|
||||||
if reg is None:
|
|
||||||
reg = _default_registry
|
|
||||||
|
|
||||||
context = request.context.model_dump()
|
|
||||||
agent_name = await classify_intent(request.message, context, reg)
|
|
||||||
response_text = await reg.call_agent(agent_name, request.message, context)
|
|
||||||
|
|
||||||
chunk_size = 50
|
|
||||||
for i in range(0, len(response_text), chunk_size):
|
|
||||||
yield response_text[i : i + chunk_size]
|
|
||||||
71
app/core/output_formatter.py
Normal file
71
app/core/output_formatter.py
Normal file
@@ -0,0 +1,71 @@
|
|||||||
|
"""Output formatter for deep-agent stream events."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from app.schemas import WsFloatingDomain, WsStreamEnd, WsStreamStart, WsStreamText
|
||||||
|
|
||||||
|
# Matches <canvas kind="...">...</canvas> blocks (single-line or multiline).
|
||||||
|
_CANVAS_BLOCK_RE = re.compile(
|
||||||
|
r'<canvas\s+kind=["\']([^"\']+)["\']>(.*?)</canvas>',
|
||||||
|
re.DOTALL | re.IGNORECASE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_canvas_block(text: str) -> tuple[str, str | None, str | None]:
|
||||||
|
"""Strip the first <canvas kind="...">...</canvas> block from *text*.
|
||||||
|
|
||||||
|
Returns ``(visible_text, canvas_content, canvas_kind)``.
|
||||||
|
``canvas_content`` and ``canvas_kind`` are ``None`` when no block is found.
|
||||||
|
"""
|
||||||
|
match = _CANVAS_BLOCK_RE.search(text)
|
||||||
|
if not match:
|
||||||
|
return text, None, None
|
||||||
|
|
||||||
|
canvas_kind = match.group(1).strip()
|
||||||
|
canvas_content = match.group(2).strip()
|
||||||
|
visible = text[: match.start()] + text[match.end() :]
|
||||||
|
visible = visible.strip()
|
||||||
|
return visible, canvas_content, canvas_kind
|
||||||
|
|
||||||
|
WsFrame = WsStreamStart | WsStreamText | WsStreamEnd | WsFloatingDomain
|
||||||
|
|
||||||
|
|
||||||
|
class StreamFormatter:
|
||||||
|
"""Convert `(event_type, data)` stream events into websocket frame models."""
|
||||||
|
|
||||||
|
def __init__(self, request_id: str) -> None:
|
||||||
|
self.request_id = request_id
|
||||||
|
|
||||||
|
async def format(
|
||||||
|
self,
|
||||||
|
event_stream: AsyncGenerator[tuple[str, Any], None],
|
||||||
|
) -> AsyncGenerator[WsFrame, None]:
|
||||||
|
started = False
|
||||||
|
|
||||||
|
async for event_type, data in event_stream:
|
||||||
|
if event_type == "floating_domain":
|
||||||
|
if isinstance(data, dict):
|
||||||
|
yield WsFloatingDomain(
|
||||||
|
request_id=self.request_id,
|
||||||
|
domain=data,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
|
if event_type != "token":
|
||||||
|
continue
|
||||||
|
|
||||||
|
if not started:
|
||||||
|
yield WsStreamStart(request_id=self.request_id)
|
||||||
|
started = True
|
||||||
|
|
||||||
|
text = str(data or "")
|
||||||
|
if text:
|
||||||
|
yield WsStreamText(request_id=self.request_id, chunk=text)
|
||||||
|
|
||||||
|
if not started:
|
||||||
|
yield WsStreamStart(request_id=self.request_id)
|
||||||
|
yield WsStreamEnd(request_id=self.request_id)
|
||||||
104
app/core/preprocessors/__init__.py
Normal file
104
app/core/preprocessors/__init__.py
Normal file
@@ -0,0 +1,104 @@
|
|||||||
|
"""Preprocessor registry: detect content type and dispatch to handlers.
|
||||||
|
|
||||||
|
Public API
|
||||||
|
----------
|
||||||
|
detect_content_type(filename, raw_content) -> str
|
||||||
|
Heuristic detection based on file extension and content patterns.
|
||||||
|
|
||||||
|
preprocess(content_type, raw_content) -> PreprocessResult
|
||||||
|
Dispatch to the appropriate handler.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
|
||||||
|
from app.core.preprocessors.base import PreprocessResult
|
||||||
|
|
||||||
|
# ── Heuristics ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
# Patterns that strongly suggest an email HTML file
|
||||||
|
_EMAIL_SIGNALS = re.compile(
|
||||||
|
r"(Subject:|From:|To:|Date:|Sent:|MIME-Version:|Content-Type:\s*text/html)",
|
||||||
|
re.IGNORECASE,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Patterns that suggest a generic HTML page (not an email)
|
||||||
|
_GENERIC_HTML_SIGNALS = re.compile(
|
||||||
|
r"<(nav|main|header|footer|article|section)\b",
|
||||||
|
re.IGNORECASE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def detect_content_type(filename: str, raw_content: str) -> str:
|
||||||
|
"""Return a content-type string for the given file.
|
||||||
|
|
||||||
|
Supported types: ``"email_html"``, ``"generic_html"``,
|
||||||
|
``"plain_text"``, ``"unknown"``.
|
||||||
|
"""
|
||||||
|
ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
|
||||||
|
|
||||||
|
if ext == "txt":
|
||||||
|
return "plain_text"
|
||||||
|
|
||||||
|
if ext in ("html", "htm", "eml", "mhtml", "mht"):
|
||||||
|
# Prefer email detection over generic HTML
|
||||||
|
if _EMAIL_SIGNALS.search(raw_content[:4096]):
|
||||||
|
return "email_html"
|
||||||
|
if _GENERIC_HTML_SIGNALS.search(raw_content[:4096]) or "<html" in raw_content[:200].lower():
|
||||||
|
return "generic_html"
|
||||||
|
# .html without clear signals — check for any email header
|
||||||
|
if re.search(r"^(From|To|Subject|Date):", raw_content[:2048], re.MULTILINE | re.IGNORECASE):
|
||||||
|
return "email_html"
|
||||||
|
return "generic_html"
|
||||||
|
|
||||||
|
# Plain text files with email headers
|
||||||
|
if ext in ("", "txt") or not ext:
|
||||||
|
if _EMAIL_SIGNALS.search(raw_content[:4096]):
|
||||||
|
return "email_html"
|
||||||
|
|
||||||
|
# Detect binary content
|
||||||
|
try:
|
||||||
|
raw_content.encode("utf-8")
|
||||||
|
except (UnicodeEncodeError, AttributeError):
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
# Non-text bytes heuristic: high ratio of non-printable chars
|
||||||
|
sample = raw_content[:512]
|
||||||
|
non_printable = sum(1 for c in sample if ord(c) < 32 and c not in "\r\n\t")
|
||||||
|
if len(sample) > 0 and non_printable / len(sample) > 0.1:
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
|
||||||
|
# ── Generic fallback handler ──────────────────────────────────────────
|
||||||
|
|
||||||
|
def _preprocess_generic(raw_content: str, content_type: str) -> PreprocessResult:
|
||||||
|
"""Strip HTML tags if present, return text as-is."""
|
||||||
|
try:
|
||||||
|
from bs4 import BeautifulSoup
|
||||||
|
text = BeautifulSoup(raw_content, "html.parser").get_text(separator="\n")
|
||||||
|
except ImportError:
|
||||||
|
# No BeautifulSoup — strip tags with a simple regex
|
||||||
|
text = re.sub(r"<[^>]+>", "", raw_content)
|
||||||
|
|
||||||
|
text = re.sub(r"\n{3,}", "\n\n", text).strip()
|
||||||
|
return PreprocessResult(content_type=content_type, clean_text=text, metadata={})
|
||||||
|
|
||||||
|
|
||||||
|
# ── Dispatch ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def preprocess(content_type: str, raw_content: str) -> PreprocessResult:
|
||||||
|
"""Dispatch *raw_content* to the handler registered for *content_type*.
|
||||||
|
|
||||||
|
Falls back to the generic handler for unknown types.
|
||||||
|
"""
|
||||||
|
if content_type == "email_html":
|
||||||
|
from app.core.preprocessors.email_html import preprocess_email_html
|
||||||
|
return preprocess_email_html(raw_content)
|
||||||
|
|
||||||
|
return _preprocess_generic(raw_content, content_type)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["detect_content_type", "preprocess", "PreprocessResult"]
|
||||||
25
app/core/preprocessors/base.py
Normal file
25
app/core/preprocessors/base.py
Normal file
@@ -0,0 +1,25 @@
|
|||||||
|
"""Base types for the preprocessor system."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class PreprocessResult:
|
||||||
|
"""Output of a preprocessor handler.
|
||||||
|
|
||||||
|
Attributes
|
||||||
|
----------
|
||||||
|
content_type:
|
||||||
|
The detected content type (e.g. ``"email_html"``, ``"plain_text"``).
|
||||||
|
clean_text:
|
||||||
|
Human-readable text stripped of markup/binary noise.
|
||||||
|
metadata:
|
||||||
|
Dict of extracted metadata (keys vary by handler).
|
||||||
|
Common keys: ``subject``, ``from``, ``to``, ``date``, ``filename``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
content_type: str
|
||||||
|
clean_text: str
|
||||||
|
metadata: dict = field(default_factory=dict)
|
||||||
111
app/core/preprocessors/email_html.py
Normal file
111
app/core/preprocessors/email_html.py
Normal file
@@ -0,0 +1,111 @@
|
|||||||
|
"""Preprocessor for email HTML files.
|
||||||
|
|
||||||
|
Handles:
|
||||||
|
- HTML stripping via BeautifulSoup
|
||||||
|
- Metadata extraction (Subject, From, To, Date)
|
||||||
|
- Thread splitting — isolates the latest reply
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from app.core.preprocessors.base import PreprocessResult
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# ── Thread split markers ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
# Matches patterns like:
|
||||||
|
# "On Mon, Apr 7, 2026 at 10:00 AM, Alice <alice@co.com> wrote:"
|
||||||
|
# "-----Original Message-----"
|
||||||
|
# "> " (plain-text quote prefix)
|
||||||
|
_THREAD_PATTERNS = [
|
||||||
|
re.compile(r"^On\s+.+wrote\s*:", re.IGNORECASE | re.MULTILINE),
|
||||||
|
re.compile(r"^-{3,}\s*(original message|forwarded message)\s*-{3,}", re.IGNORECASE | re.MULTILINE),
|
||||||
|
re.compile(r"^>{1,}\s+\S", re.MULTILINE),
|
||||||
|
re.compile(r"^From:\s+.+\nSent:\s+", re.IGNORECASE | re.MULTILINE),
|
||||||
|
]
|
||||||
|
|
||||||
|
# ── Metadata patterns (applied on raw HTML / plain fallback) ──────────
|
||||||
|
|
||||||
|
_META_PATTERNS: dict[str, list[re.Pattern]] = {
|
||||||
|
"subject": [
|
||||||
|
re.compile(r"<title>(.+?)</title>", re.IGNORECASE | re.DOTALL),
|
||||||
|
re.compile(r"Subject:\s*(.+)", re.IGNORECASE),
|
||||||
|
],
|
||||||
|
"from": [
|
||||||
|
re.compile(r'<meta[^>]+name=["\']?from["\']?[^>]+content=["\']([^"\']+)["\']', re.IGNORECASE),
|
||||||
|
re.compile(r"From:\s*(.+)", re.IGNORECASE),
|
||||||
|
],
|
||||||
|
"to": [
|
||||||
|
re.compile(r'<meta[^>]+name=["\']?to["\']?[^>]+content=["\']([^"\']+)["\']', re.IGNORECASE),
|
||||||
|
re.compile(r"To:\s*(.+)", re.IGNORECASE),
|
||||||
|
],
|
||||||
|
"date": [
|
||||||
|
re.compile(r'<meta[^>]+name=["\']?date["\']?[^>]+content=["\']([^"\']+)["\']', re.IGNORECASE),
|
||||||
|
re.compile(r"Date:\s*(.+)", re.IGNORECASE),
|
||||||
|
re.compile(r"Sent:\s*(.+)", re.IGNORECASE),
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_metadata(raw_html: str, text: str) -> dict:
|
||||||
|
"""Extract Subject/From/To/Date from raw HTML or plain text."""
|
||||||
|
metadata: dict[str, str] = {}
|
||||||
|
for field, patterns in _META_PATTERNS.items():
|
||||||
|
for pat in patterns:
|
||||||
|
m = pat.search(raw_html) or pat.search(text)
|
||||||
|
if m:
|
||||||
|
metadata[field] = m.group(1).strip()
|
||||||
|
break
|
||||||
|
return metadata
|
||||||
|
|
||||||
|
|
||||||
|
def _split_thread(text: str) -> str:
|
||||||
|
"""Return only the latest message in a threaded email."""
|
||||||
|
earliest_pos: int | None = None
|
||||||
|
for pat in _THREAD_PATTERNS:
|
||||||
|
m = pat.search(text)
|
||||||
|
if m and (earliest_pos is None or m.start() < earliest_pos):
|
||||||
|
earliest_pos = m.start()
|
||||||
|
|
||||||
|
if earliest_pos is not None and earliest_pos > 0:
|
||||||
|
return text[:earliest_pos].strip()
|
||||||
|
return text.strip()
|
||||||
|
|
||||||
|
|
||||||
|
def preprocess_email_html(raw_content: str) -> PreprocessResult:
|
||||||
|
"""Strip HTML, extract metadata, split thread from an email HTML file."""
|
||||||
|
try:
|
||||||
|
from bs4 import BeautifulSoup # lazy import — optional dep
|
||||||
|
except ImportError as exc:
|
||||||
|
raise ImportError(
|
||||||
|
"beautifulsoup4 is required for email_html preprocessing. "
|
||||||
|
"Install it with: pip install beautifulsoup4"
|
||||||
|
) from exc
|
||||||
|
|
||||||
|
# Parse with lxml if available, fall back to html.parser
|
||||||
|
try:
|
||||||
|
soup = BeautifulSoup(raw_content, "lxml")
|
||||||
|
except Exception:
|
||||||
|
soup = BeautifulSoup(raw_content, "html.parser")
|
||||||
|
|
||||||
|
# Remove noise tags
|
||||||
|
for tag in soup(["style", "script", "head", "noscript"]):
|
||||||
|
tag.decompose()
|
||||||
|
|
||||||
|
clean_text = soup.get_text(separator="\n")
|
||||||
|
# Collapse excessive blank lines
|
||||||
|
clean_text = re.sub(r"\n{3,}", "\n\n", clean_text).strip()
|
||||||
|
|
||||||
|
metadata = _extract_metadata(raw_content, clean_text)
|
||||||
|
latest_message = _split_thread(clean_text)
|
||||||
|
|
||||||
|
return PreprocessResult(
|
||||||
|
content_type="email_html",
|
||||||
|
clean_text=latest_message,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
@@ -7,16 +7,54 @@ The callback sends a `tool_call` WS frame and awaits the `tool_result`.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
from contextvars import ContextVar
|
from contextvars import ContextVar
|
||||||
from typing import Any, Callable, Coroutine
|
from typing import Any, Callable, Coroutine
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
|
_SNAKE_TO_CAMEL_RE = re.compile(r"_([a-z])")
|
||||||
|
|
||||||
|
|
||||||
|
def _key_to_camel(key: str) -> str:
|
||||||
|
return _SNAKE_TO_CAMEL_RE.sub(lambda m: m.group(1).upper(), key)
|
||||||
|
|
||||||
|
|
||||||
|
def _keys_to_camel(obj: Any) -> Any:
|
||||||
|
"""Recursively convert dict keys from snake_case to camelCase.
|
||||||
|
|
||||||
|
Mirrors the JS-side ``toCamelCase`` applied to incoming WS frames in
|
||||||
|
``adiuvAI/src/main/api/backend-client.ts``. The Electron executor wraps
|
||||||
|
tool_result payloads in ``toSnakeCase`` before sending; this restores the
|
||||||
|
camelCase schema property names that the tool code expects to read.
|
||||||
|
"""
|
||||||
|
if isinstance(obj, dict):
|
||||||
|
return {_key_to_camel(k): _keys_to_camel(v) for k, v in obj.items()}
|
||||||
|
if isinstance(obj, list):
|
||||||
|
return [_keys_to_camel(v) for v in obj]
|
||||||
|
return obj
|
||||||
|
|
||||||
# Holds the execute callback for the current WS session.
|
# Holds the execute callback for the current WS session.
|
||||||
# Set by the chat WS handler before the orchestrator runs; cleared after.
|
# Set by the chat WS handler before the orchestrator runs; cleared after.
|
||||||
_client_executor: ContextVar[Callable[[dict], Coroutine[Any, Any, dict]]] = ContextVar(
|
_client_executor: ContextVar[Callable[[dict], Coroutine[Any, Any, dict]]] = ContextVar(
|
||||||
"_client_executor"
|
"_client_executor"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Optional collector that captures raw execute_on_client results.
|
||||||
|
# Set by _tool_loop / _tool_loop_stream to populate ChatAgent.tool_results.
|
||||||
|
_tool_result_collector: ContextVar[list[dict] | None] = ContextVar(
|
||||||
|
"_tool_result_collector", default=None
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def set_tool_result_collector(lst: list[dict]) -> None:
|
||||||
|
"""Register *lst* as the collector for this async context."""
|
||||||
|
_tool_result_collector.set(lst)
|
||||||
|
|
||||||
|
|
||||||
|
def clear_tool_result_collector() -> None:
|
||||||
|
"""Clear the collector (best-effort)."""
|
||||||
|
_tool_result_collector.set(None)
|
||||||
|
|
||||||
|
|
||||||
def set_client_executor(fn: Callable[[dict], Coroutine[Any, Any, dict]]) -> None:
|
def set_client_executor(fn: Callable[[dict], Coroutine[Any, Any, dict]]) -> None:
|
||||||
"""Bind *fn* as the executor for the current async context (task/coroutine)."""
|
"""Bind *fn* as the executor for the current async context (task/coroutine)."""
|
||||||
@@ -65,4 +103,13 @@ async def execute_on_client(
|
|||||||
if limit is not None:
|
if limit is not None:
|
||||||
payload["limit"] = limit
|
payload["limit"] = limit
|
||||||
|
|
||||||
return await callback(payload)
|
result = await callback(payload)
|
||||||
|
result = _keys_to_camel(result)
|
||||||
|
collector = _tool_result_collector.get(None)
|
||||||
|
if collector is not None:
|
||||||
|
collector.append({
|
||||||
|
"action": action,
|
||||||
|
"table": table,
|
||||||
|
"data": result,
|
||||||
|
})
|
||||||
|
return result
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ from app.config.settings import settings
|
|||||||
engine = create_async_engine(
|
engine = create_async_engine(
|
||||||
settings.DATABASE_URL,
|
settings.DATABASE_URL,
|
||||||
pool_pre_ping=True,
|
pool_pre_ping=True,
|
||||||
echo=settings.ENV == "dev",
|
echo=False,
|
||||||
)
|
)
|
||||||
|
|
||||||
async_session = async_sessionmaker(engine, expire_on_commit=False)
|
async_session = async_sessionmaker(engine, expire_on_commit=False)
|
||||||
|
|||||||
164
app/integrations/__init__.py
Normal file
164
app/integrations/__init__.py
Normal file
@@ -0,0 +1,164 @@
|
|||||||
|
"""Cloud provider integration utilities.
|
||||||
|
|
||||||
|
Provides:
|
||||||
|
* Shared message dataclasses (``EmailMessage``, ``ChatMessage``) used by
|
||||||
|
both the Gmail and MS Graph clients and consumed by ``agent_runner``.
|
||||||
|
* ``get_provider()`` — factory that returns the correct client given a
|
||||||
|
provider name and decrypted OAuth credentials dict.
|
||||||
|
* ``encrypt_token()`` / ``decrypt_token()`` — Fernet-based at-rest
|
||||||
|
encryption for OAuth tokens stored in ``cloud_agent_configs``.
|
||||||
|
|
||||||
|
Encryption rationale
|
||||||
|
--------------------
|
||||||
|
Unlike user content (which is E2E-encrypted client-side and **never**
|
||||||
|
decrypted server-side), OAuth tokens *must* be decrypted server-side
|
||||||
|
because the backend makes provider API calls on behalf of the user.
|
||||||
|
The Fernet key lives solely in ``OAUTH_ENCRYPTION_KEY`` env var — it
|
||||||
|
is never returned to clients.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from cryptography.fernet import Fernet, InvalidToken
|
||||||
|
|
||||||
|
from app.config.settings import settings
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from app.integrations.gmail import GmailClient
|
||||||
|
from app.integrations.ms_graph import MSGraphClient
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ── Shared message types ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class EmailMessage:
|
||||||
|
"""A single email message fetched from Gmail or Outlook."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
subject: str
|
||||||
|
sender: str
|
||||||
|
body_text: str
|
||||||
|
date: datetime
|
||||||
|
labels: list[str] = field(default_factory=list)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def as_text(self) -> str:
|
||||||
|
"""Return a human-readable text representation for LLM extraction."""
|
||||||
|
date_str = self.date.strftime("%Y-%m-%d %H:%M")
|
||||||
|
labels_str = f" [{', '.join(self.labels)}]" if self.labels else ""
|
||||||
|
return (
|
||||||
|
f"From: {self.sender}\n"
|
||||||
|
f"Date: {date_str}{labels_str}\n"
|
||||||
|
f"Subject: {self.subject}\n\n"
|
||||||
|
f"{self.body_text}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ChatMessage:
|
||||||
|
"""A single Teams chat or channel message fetched from MS Graph."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
content: str
|
||||||
|
sender: str
|
||||||
|
channel: str | None
|
||||||
|
date: datetime
|
||||||
|
|
||||||
|
@property
|
||||||
|
def as_text(self) -> str:
|
||||||
|
"""Return a human-readable text representation for LLM extraction."""
|
||||||
|
date_str = self.date.strftime("%Y-%m-%d %H:%M")
|
||||||
|
channel_str = f" [channel: {self.channel}]" if self.channel else ""
|
||||||
|
return (
|
||||||
|
f"From: {self.sender}\n"
|
||||||
|
f"Date: {date_str}{channel_str}\n\n"
|
||||||
|
f"{self.content}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Fernet helpers ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def _get_fernet() -> Fernet:
|
||||||
|
"""Return a ``Fernet`` instance using ``settings.OAUTH_ENCRYPTION_KEY``.
|
||||||
|
|
||||||
|
Raises ``RuntimeError`` if ``OAUTH_ENCRYPTION_KEY`` is not set — callers
|
||||||
|
must ensure this is configured before persisting OAuth tokens.
|
||||||
|
"""
|
||||||
|
key = settings.OAUTH_ENCRYPTION_KEY
|
||||||
|
if not key:
|
||||||
|
raise RuntimeError(
|
||||||
|
"OAUTH_ENCRYPTION_KEY is not set. "
|
||||||
|
"Generate one with: python -c \"from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())\""
|
||||||
|
)
|
||||||
|
return Fernet(key.encode() if isinstance(key, str) else key)
|
||||||
|
|
||||||
|
|
||||||
|
def encrypt_token(token_info: dict) -> str:
|
||||||
|
"""Fernet-encrypt an OAuth credential dict and return a base64 string.
|
||||||
|
|
||||||
|
Stores the full ``{access_token, refresh_token, token_uri, client_id,
|
||||||
|
client_secret, scopes, expiry}`` dict (or equivalent MSAL shape).
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: OAUTH_ENCRYPTION_KEY is not configured.
|
||||||
|
ValueError: ``token_info`` is not a non-empty dict.
|
||||||
|
"""
|
||||||
|
if not isinstance(token_info, dict) or not token_info:
|
||||||
|
raise ValueError("token_info must be a non-empty dict")
|
||||||
|
plaintext = json.dumps(token_info).encode("utf-8")
|
||||||
|
return _get_fernet().encrypt(plaintext).decode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
def decrypt_token(encrypted: str) -> dict:
|
||||||
|
"""Decrypt a Fernet-encrypted token string and return the credential dict.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: OAUTH_ENCRYPTION_KEY is not configured.
|
||||||
|
ValueError: The encrypted string is invalid or was encrypted with a
|
||||||
|
different key.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
plaintext = _get_fernet().decrypt(encrypted.encode("utf-8"))
|
||||||
|
return json.loads(plaintext)
|
||||||
|
except (InvalidToken, json.JSONDecodeError) as exc:
|
||||||
|
raise ValueError(f"Failed to decrypt OAuth token: {exc}") from exc
|
||||||
|
|
||||||
|
|
||||||
|
# ── Provider factory ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def get_provider(
|
||||||
|
provider: str,
|
||||||
|
credentials_info: dict,
|
||||||
|
) -> "GmailClient | MSGraphClient":
|
||||||
|
"""Return the correct provider client for *provider*.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
provider:
|
||||||
|
One of ``"gmail"``, ``"outlook"``, ``"teams"``.
|
||||||
|
credentials_info:
|
||||||
|
Decrypted OAuth credential dict (Google or Microsoft shape).
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: Unknown provider name.
|
||||||
|
"""
|
||||||
|
if provider == "gmail":
|
||||||
|
from app.integrations.gmail import GmailClient
|
||||||
|
return GmailClient(credentials_info)
|
||||||
|
if provider in {"outlook", "teams"}:
|
||||||
|
from app.integrations.ms_graph import MSGraphClient
|
||||||
|
return MSGraphClient(credentials_info)
|
||||||
|
raise ValueError(
|
||||||
|
f"Unknown cloud provider {provider!r}. "
|
||||||
|
"Supported: 'gmail', 'outlook', 'teams'."
|
||||||
|
)
|
||||||
335
app/integrations/gmail.py
Normal file
335
app/integrations/gmail.py
Normal file
@@ -0,0 +1,335 @@
|
|||||||
|
"""Gmail API client for cloud agent integration.
|
||||||
|
|
||||||
|
Wraps the Google Gmail REST API to fetch email messages matching a
|
||||||
|
``filter_config`` dict. Uses the official ``google-api-python-client``
|
||||||
|
library (synchronous) wrapped in ``asyncio.to_thread()`` to avoid
|
||||||
|
blocking the event loop.
|
||||||
|
|
||||||
|
Token refresh is handled transparently: when the stored access token has
|
||||||
|
expired, ``google.auth.transport.requests.Request`` will use the refresh
|
||||||
|
token to obtain a fresh one. The caller is responsible for persisting
|
||||||
|
any refreshed credentials back to ``CloudAgentConfig.oauth_token_encrypted``
|
||||||
|
(see ``agent_runner.run_cloud_agent``).
|
||||||
|
|
||||||
|
Credential dict shape (Google OAuth2):
|
||||||
|
{
|
||||||
|
"token": "<access_token>",
|
||||||
|
"refresh_token": "<refresh_token>",
|
||||||
|
"token_uri": "https://oauth2.googleapis.com/token",
|
||||||
|
"client_id": "<client_id>",
|
||||||
|
"client_secret": "<client_secret>",
|
||||||
|
"scopes": ["https://www.googleapis.com/auth/gmail.readonly"],
|
||||||
|
"expiry": "2025-01-01T00:00:00Z" # optional ISO-8601
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import email
|
||||||
|
import html
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from app.integrations import EmailMessage
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Gmail search date format — e.g. "after:2025/01/01"
|
||||||
|
_GMAIL_DATE_FMT = "%Y/%m/%d"
|
||||||
|
|
||||||
|
# Maximum characters of body text forwarded to the LLM.
|
||||||
|
_BODY_TRUNCATE = 8_000
|
||||||
|
|
||||||
|
# Maximum messages retrieved per run (prevents runaway quota usage).
|
||||||
|
_MAX_MESSAGES = 200
|
||||||
|
|
||||||
|
|
||||||
|
def _build_gmail_query(
|
||||||
|
filter_config: dict[str, Any] | None,
|
||||||
|
since: datetime | None,
|
||||||
|
) -> str:
|
||||||
|
"""Build a Gmail search query string from *filter_config* and *since*.
|
||||||
|
|
||||||
|
Supported ``filter_config`` keys:
|
||||||
|
labels (list[str]): Gmail label names, e.g. ``["INBOX", "work"]``
|
||||||
|
senders (list[str]): Sender addresses or domains to include
|
||||||
|
date_range (dict): ``{from: "<YYYY-MM-DD>", to: "<YYYY-MM-DD>"}``
|
||||||
|
|
||||||
|
A hard ``since`` date (from last run) always overrides ``date_range.from``
|
||||||
|
when it is earlier.
|
||||||
|
"""
|
||||||
|
parts: list[str] = []
|
||||||
|
cfg = filter_config or {}
|
||||||
|
|
||||||
|
# Labels — joined with OR when multiple given.
|
||||||
|
labels: list[str] = cfg.get("labels", [])
|
||||||
|
if labels:
|
||||||
|
if len(labels) == 1:
|
||||||
|
parts.append(f"label:{labels[0]}")
|
||||||
|
else:
|
||||||
|
label_expr = " OR ".join(f"label:{lbl}" for lbl in labels)
|
||||||
|
parts.append(f"({label_expr})")
|
||||||
|
|
||||||
|
# Senders — each prefixed with "from:".
|
||||||
|
senders: list[str] = cfg.get("senders", [])
|
||||||
|
for sender in senders:
|
||||||
|
parts.append(f"from:{sender}")
|
||||||
|
|
||||||
|
# Date range.
|
||||||
|
date_range: dict = cfg.get("date_range", {})
|
||||||
|
from_str: str | None = date_range.get("from")
|
||||||
|
to_str: str | None = date_range.get("to")
|
||||||
|
|
||||||
|
# Determine effective "from" date: most recent of filter_config.date_range.from and since.
|
||||||
|
effective_since: datetime | None = since
|
||||||
|
if from_str:
|
||||||
|
try:
|
||||||
|
cfg_since = datetime.fromisoformat(from_str.replace("Z", "+00:00"))
|
||||||
|
if cfg_since.tzinfo is None:
|
||||||
|
cfg_since = cfg_since.replace(tzinfo=timezone.utc)
|
||||||
|
if effective_since is None or cfg_since > effective_since:
|
||||||
|
effective_since = cfg_since
|
||||||
|
except ValueError:
|
||||||
|
logger.warning("gmail: invalid date_range.from %r — ignoring", from_str)
|
||||||
|
|
||||||
|
if effective_since:
|
||||||
|
parts.append(f"after:{effective_since.strftime(_GMAIL_DATE_FMT)}")
|
||||||
|
|
||||||
|
if to_str:
|
||||||
|
try:
|
||||||
|
to_dt = datetime.fromisoformat(to_str.replace("Z", "+00:00"))
|
||||||
|
parts.append(f"before:{to_dt.strftime(_GMAIL_DATE_FMT)}")
|
||||||
|
except ValueError:
|
||||||
|
logger.warning("gmail: invalid date_range.to %r — ignoring", to_str)
|
||||||
|
|
||||||
|
return " ".join(parts)
|
||||||
|
|
||||||
|
|
||||||
|
def _strip_html(raw_html: str) -> str:
|
||||||
|
"""Remove HTML tags and decode entities to get plain text."""
|
||||||
|
no_tags = re.sub(r"<[^>]+>", " ", raw_html)
|
||||||
|
decoded = html.unescape(no_tags)
|
||||||
|
return re.sub(r"\s+", " ", decoded).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_body(payload: dict[str, Any]) -> str:
|
||||||
|
"""Recursively extract the plain-text body from a Gmail message payload.
|
||||||
|
|
||||||
|
Prefers ``text/plain``; falls back to ``text/html`` (stripped of tags).
|
||||||
|
Returns an empty string if no body can be extracted.
|
||||||
|
"""
|
||||||
|
mime_type: str = payload.get("mimeType", "")
|
||||||
|
body: dict = payload.get("body", {})
|
||||||
|
parts: list[dict] = payload.get("parts", [])
|
||||||
|
|
||||||
|
if mime_type == "text/plain":
|
||||||
|
data = body.get("data", "")
|
||||||
|
if data:
|
||||||
|
return base64.urlsafe_b64decode(data + "==").decode("utf-8", errors="replace")
|
||||||
|
return ""
|
||||||
|
|
||||||
|
if mime_type == "text/html":
|
||||||
|
data = body.get("data", "")
|
||||||
|
if data:
|
||||||
|
raw = base64.urlsafe_b64decode(data + "==").decode("utf-8", errors="replace")
|
||||||
|
return _strip_html(raw)
|
||||||
|
return ""
|
||||||
|
|
||||||
|
# Multipart — prefer text/plain part, fall back to text/html.
|
||||||
|
plain_fallback = ""
|
||||||
|
for part in parts:
|
||||||
|
part_mime = part.get("mimeType", "")
|
||||||
|
if part_mime == "text/plain":
|
||||||
|
return _parse_body(part)
|
||||||
|
if part_mime == "text/html" and not plain_fallback:
|
||||||
|
plain_fallback = _parse_body(part)
|
||||||
|
if part_mime.startswith("multipart/"):
|
||||||
|
nested = _parse_body(part)
|
||||||
|
if nested:
|
||||||
|
return nested
|
||||||
|
return plain_fallback
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_date(raw: str) -> datetime:
|
||||||
|
"""Parse an RFC 2822 email date header into a UTC ``datetime``."""
|
||||||
|
try:
|
||||||
|
parsed = email.utils.parsedate_to_datetime(raw)
|
||||||
|
if parsed.tzinfo is None:
|
||||||
|
parsed = parsed.replace(tzinfo=timezone.utc)
|
||||||
|
return parsed.astimezone(timezone.utc)
|
||||||
|
except Exception:
|
||||||
|
return datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
|
||||||
|
class GmailClient:
|
||||||
|
"""Fetch email messages from a Gmail account via the Gmail REST API.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
credentials_info:
|
||||||
|
Decrypted OAuth2 credential dict. Must contain at minimum
|
||||||
|
``token`` (access token) or ``refresh_token`` + ``token_uri`` +
|
||||||
|
``client_id`` + ``client_secret``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, credentials_info: dict[str, Any]) -> None:
|
||||||
|
from google.oauth2.credentials import Credentials
|
||||||
|
|
||||||
|
self._credentials_info = credentials_info
|
||||||
|
expiry_str: str | None = credentials_info.get("expiry")
|
||||||
|
expiry: datetime | None = None
|
||||||
|
if expiry_str:
|
||||||
|
try:
|
||||||
|
expiry = datetime.fromisoformat(
|
||||||
|
expiry_str.replace("Z", "+00:00")
|
||||||
|
).replace(tzinfo=timezone.utc)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
self._credentials = Credentials(
|
||||||
|
token=credentials_info.get("token"),
|
||||||
|
refresh_token=credentials_info.get("refresh_token"),
|
||||||
|
token_uri=credentials_info.get("token_uri", "https://oauth2.googleapis.com/token"),
|
||||||
|
client_id=credentials_info.get("client_id"),
|
||||||
|
client_secret=credentials_info.get("client_secret"),
|
||||||
|
scopes=credentials_info.get("scopes"),
|
||||||
|
expiry=expiry,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Public API ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def fetch_messages(
|
||||||
|
self,
|
||||||
|
filter_config: dict[str, Any] | None = None,
|
||||||
|
since: datetime | None = None,
|
||||||
|
) -> list[EmailMessage]:
|
||||||
|
"""Return up to ``_MAX_MESSAGES`` emails matching *filter_config*.
|
||||||
|
|
||||||
|
Runs the synchronous Google API calls inside ``asyncio.to_thread()``
|
||||||
|
to avoid blocking the async event loop.
|
||||||
|
|
||||||
|
Token refresh is performed automatically when the access token has
|
||||||
|
expired. After the call, ``self.refreshed_credentials`` may be
|
||||||
|
consulted to detect whether new credentials should be persisted.
|
||||||
|
"""
|
||||||
|
query = _build_gmail_query(filter_config, since)
|
||||||
|
logger.debug("gmail: executing search query %r", query)
|
||||||
|
return await asyncio.to_thread(self._fetch_sync, query)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def refreshed_credentials(self) -> dict[str, Any] | None:
|
||||||
|
"""Return updated credential dict if the access token was refreshed.
|
||||||
|
|
||||||
|
If the credentials were refreshed during ``fetch_messages()``, returns
|
||||||
|
a new dict that should be re-encrypted and written back to the DB.
|
||||||
|
Returns ``None`` if no refresh occurred.
|
||||||
|
"""
|
||||||
|
creds = self._credentials
|
||||||
|
if not creds.valid and creds.expired:
|
||||||
|
return None
|
||||||
|
# Check whether the token changed from what was stored.
|
||||||
|
if creds.token != self._credentials_info.get("token"):
|
||||||
|
result = {
|
||||||
|
"token": creds.token,
|
||||||
|
"refresh_token": creds.refresh_token,
|
||||||
|
"token_uri": creds.token_uri,
|
||||||
|
"client_id": creds.client_id,
|
||||||
|
"client_secret": creds.client_secret,
|
||||||
|
"scopes": list(creds.scopes or []),
|
||||||
|
}
|
||||||
|
if creds.expiry:
|
||||||
|
result["expiry"] = creds.expiry.isoformat()
|
||||||
|
return result
|
||||||
|
return None
|
||||||
|
|
||||||
|
# ── Internal sync worker ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
def _fetch_sync(self, query: str) -> list[EmailMessage]:
|
||||||
|
"""Synchronous worker — called inside ``asyncio.to_thread()``."""
|
||||||
|
import googleapiclient.discovery
|
||||||
|
import googleapiclient.errors
|
||||||
|
from google.auth.transport.requests import Request
|
||||||
|
|
||||||
|
# Refresh token if needed before building the service.
|
||||||
|
if self._credentials.expired and self._credentials.refresh_token:
|
||||||
|
try:
|
||||||
|
self._credentials.refresh(Request())
|
||||||
|
except Exception as exc:
|
||||||
|
raise RuntimeError(f"Gmail token refresh failed: {exc}") from exc
|
||||||
|
|
||||||
|
service = googleapiclient.discovery.build(
|
||||||
|
"gmail", "v1", credentials=self._credentials, cache_discovery=False
|
||||||
|
)
|
||||||
|
user_api = service.users() # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
# ── List matching message IDs ──────────────────────────────────────
|
||||||
|
ids: list[str] = []
|
||||||
|
page_token: str | None = None
|
||||||
|
while len(ids) < _MAX_MESSAGES:
|
||||||
|
batch_size = min(100, _MAX_MESSAGES - len(ids))
|
||||||
|
kwargs: dict[str, Any] = {
|
||||||
|
"userId": "me",
|
||||||
|
"maxResults": batch_size,
|
||||||
|
}
|
||||||
|
if query:
|
||||||
|
kwargs["q"] = query
|
||||||
|
if page_token:
|
||||||
|
kwargs["pageToken"] = page_token
|
||||||
|
|
||||||
|
try:
|
||||||
|
resp = user_api.messages().list(**kwargs).execute()
|
||||||
|
except googleapiclient.errors.HttpError as exc:
|
||||||
|
raise RuntimeError(f"Gmail messages.list failed: {exc}") from exc
|
||||||
|
|
||||||
|
for msg in resp.get("messages", []):
|
||||||
|
ids.append(msg["id"])
|
||||||
|
|
||||||
|
page_token = resp.get("nextPageToken")
|
||||||
|
if not page_token:
|
||||||
|
break
|
||||||
|
|
||||||
|
if not ids:
|
||||||
|
logger.debug("gmail: no messages matched query %r", query)
|
||||||
|
return []
|
||||||
|
|
||||||
|
logger.info("gmail: fetching %d message(s)", len(ids))
|
||||||
|
|
||||||
|
# ── Fetch individual message details ──────────────────────────────
|
||||||
|
messages: list[EmailMessage] = []
|
||||||
|
for msg_id in ids:
|
||||||
|
try:
|
||||||
|
msg = user_api.messages().get(
|
||||||
|
userId="me", id=msg_id, format="full"
|
||||||
|
).execute()
|
||||||
|
|
||||||
|
headers: dict[str, str] = {
|
||||||
|
h["name"].lower(): h["value"]
|
||||||
|
for h in msg.get("payload", {}).get("headers", [])
|
||||||
|
}
|
||||||
|
subject = headers.get("subject", "(no subject)")
|
||||||
|
sender = headers.get("from", "unknown")
|
||||||
|
date_raw = headers.get("date", "")
|
||||||
|
date = _parse_date(date_raw) if date_raw else datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
body_text = _parse_body(msg.get("payload", {}))[:_BODY_TRUNCATE]
|
||||||
|
labels = msg.get("labelIds", [])
|
||||||
|
|
||||||
|
messages.append(EmailMessage(
|
||||||
|
id=msg_id,
|
||||||
|
subject=subject,
|
||||||
|
sender=sender,
|
||||||
|
body_text=body_text,
|
||||||
|
date=date,
|
||||||
|
labels=labels,
|
||||||
|
))
|
||||||
|
except googleapiclient.errors.HttpError as exc:
|
||||||
|
logger.warning("gmail: skipping message %s — HTTP error: %s", msg_id, exc)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("gmail: skipping message %s — unexpected error: %s", msg_id, exc)
|
||||||
|
|
||||||
|
logger.info("gmail: returned %d message(s)", len(messages))
|
||||||
|
return messages
|
||||||
352
app/integrations/ms_graph.py
Normal file
352
app/integrations/ms_graph.py
Normal file
@@ -0,0 +1,352 @@
|
|||||||
|
"""Microsoft Graph API client for Outlook and Teams cloud agent integration.
|
||||||
|
|
||||||
|
Handles two data sources:
|
||||||
|
|
||||||
|
* **Outlook email** (``provider="outlook"``) — ``fetch_emails()`` calls
|
||||||
|
``/me/messages`` with an OData ``$filter`` built from ``filter_config``.
|
||||||
|
* **Teams messages** (``provider="teams"``) — ``fetch_messages()`` calls
|
||||||
|
``/me/chats/getAllMessages`` filtered by date.
|
||||||
|
|
||||||
|
Authentication uses MSAL ``PublicClientApplication`` to acquire a token
|
||||||
|
from a stored refresh token. The ``httpx.AsyncClient`` (already a project
|
||||||
|
dependency) is used for all API calls.
|
||||||
|
|
||||||
|
Credential dict shape (Microsoft OAuth2 / MSAL):
|
||||||
|
{
|
||||||
|
"access_token": "<access_token>",
|
||||||
|
"refresh_token": "<refresh_token>",
|
||||||
|
"token_type": "Bearer",
|
||||||
|
"scope": "Mail.Read ChannelMessage.Read.All offline_access",
|
||||||
|
"expires_in": 3600
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
from app.config.settings import settings
|
||||||
|
from app.integrations import ChatMessage, EmailMessage
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_GRAPH_BASE = "https://graph.microsoft.com/v1.0"
|
||||||
|
|
||||||
|
# Max items fetched per run.
|
||||||
|
_MAX_EMAILS = 200
|
||||||
|
_MAX_MESSAGES = 200
|
||||||
|
|
||||||
|
# Max characters of body forwarded to the LLM.
|
||||||
|
_BODY_TRUNCATE = 8_000
|
||||||
|
|
||||||
|
|
||||||
|
def _strip_html(raw: str) -> str:
|
||||||
|
"""Strip HTML tags and collapse whitespace."""
|
||||||
|
no_tags = re.sub(r"<[^>]+>", " ", raw)
|
||||||
|
import html as _html
|
||||||
|
decoded = _html.unescape(no_tags)
|
||||||
|
return re.sub(r"\s+", " ", decoded).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _odata_datetime(dt: datetime) -> str:
|
||||||
|
"""Format a datetime as an OData datetime literal (UTC, ISO 8601)."""
|
||||||
|
utc = dt.astimezone(timezone.utc)
|
||||||
|
return utc.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||||
|
|
||||||
|
|
||||||
|
def _build_email_filter(
|
||||||
|
filter_config: dict[str, Any] | None,
|
||||||
|
since: datetime | None,
|
||||||
|
) -> str:
|
||||||
|
"""Build an OData ``$filter`` expression for the ``/me/messages`` endpoint.
|
||||||
|
|
||||||
|
Supported ``filter_config`` keys:
|
||||||
|
senders (list[str]): Sender email addresses.
|
||||||
|
date_range (dict): ``{from: "<ISO-8601>", to: "<ISO-8601>"}``
|
||||||
|
folders (list[str]): Folder display names (not directly filterable
|
||||||
|
via OData, so ignored here — callers iterate
|
||||||
|
folder IDs separately if needed; listed for
|
||||||
|
completeness).
|
||||||
|
|
||||||
|
A hard ``since`` date always overrides ``date_range.from`` when it is
|
||||||
|
earlier.
|
||||||
|
"""
|
||||||
|
clauses: list[str] = []
|
||||||
|
cfg = filter_config or {}
|
||||||
|
|
||||||
|
# Senders.
|
||||||
|
senders: list[str] = cfg.get("senders", [])
|
||||||
|
if senders:
|
||||||
|
sender_clauses = [f"from/emailAddress/address eq '{s}'" for s in senders]
|
||||||
|
clauses.append("(" + " or ".join(sender_clauses) + ")")
|
||||||
|
|
||||||
|
# Date range.
|
||||||
|
date_range: dict = cfg.get("date_range", {})
|
||||||
|
from_str: str | None = date_range.get("from")
|
||||||
|
|
||||||
|
effective_since: datetime | None = since
|
||||||
|
if from_str:
|
||||||
|
try:
|
||||||
|
cfg_since = datetime.fromisoformat(from_str.replace("Z", "+00:00"))
|
||||||
|
if cfg_since.tzinfo is None:
|
||||||
|
cfg_since = cfg_since.replace(tzinfo=timezone.utc)
|
||||||
|
if effective_since is None or cfg_since > effective_since:
|
||||||
|
effective_since = cfg_since
|
||||||
|
except ValueError:
|
||||||
|
logger.warning("ms_graph: invalid date_range.from %r — ignoring", from_str)
|
||||||
|
|
||||||
|
if effective_since:
|
||||||
|
clauses.append(f"receivedDateTime ge {_odata_datetime(effective_since)}")
|
||||||
|
|
||||||
|
to_str: str | None = date_range.get("to")
|
||||||
|
if to_str:
|
||||||
|
try:
|
||||||
|
to_dt = datetime.fromisoformat(to_str.replace("Z", "+00:00"))
|
||||||
|
if to_dt.tzinfo is None:
|
||||||
|
to_dt = to_dt.replace(tzinfo=timezone.utc)
|
||||||
|
clauses.append(f"receivedDateTime le {_odata_datetime(to_dt)}")
|
||||||
|
except ValueError:
|
||||||
|
logger.warning("ms_graph: invalid date_range.to %r — ignoring", to_str)
|
||||||
|
|
||||||
|
return " and ".join(clauses)
|
||||||
|
|
||||||
|
|
||||||
|
class MSGraphClient:
|
||||||
|
"""Fetch emails and Teams messages via the Microsoft Graph REST API.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
credentials_info:
|
||||||
|
Decrypted MSAL credential dict.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, credentials_info: dict[str, Any]) -> None:
|
||||||
|
self._credentials_info = credentials_info
|
||||||
|
self._access_token: str = credentials_info.get("access_token", "")
|
||||||
|
self._original_access_token: str = self._access_token
|
||||||
|
self._refresh_token: str | None = credentials_info.get("refresh_token")
|
||||||
|
|
||||||
|
# ── Token management ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def _auth_headers(self) -> dict[str, str]:
|
||||||
|
return {"Authorization": f"Bearer {self._access_token}"}
|
||||||
|
|
||||||
|
async def _refresh_access_token(self) -> None:
|
||||||
|
"""Use MSAL to exchange the refresh token for a fresh access token.
|
||||||
|
|
||||||
|
Updates ``self._access_token`` and ``self._credentials_info`` in-place.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: MSAL reports an auth error.
|
||||||
|
"""
|
||||||
|
import msal
|
||||||
|
|
||||||
|
app = msal.ConfidentialClientApplication(
|
||||||
|
client_id=settings.MS_CLIENT_ID,
|
||||||
|
client_credential=settings.MS_CLIENT_SECRET,
|
||||||
|
authority=f"https://login.microsoftonline.com/{settings.MS_TENANT_ID}",
|
||||||
|
)
|
||||||
|
scopes: list[str] = self._credentials_info.get("scope", "").split()
|
||||||
|
if not scopes:
|
||||||
|
scopes = ["https://graph.microsoft.com/.default"]
|
||||||
|
|
||||||
|
result = app.acquire_token_by_refresh_token(
|
||||||
|
self._refresh_token,
|
||||||
|
scopes=scopes,
|
||||||
|
)
|
||||||
|
if "access_token" not in result:
|
||||||
|
error = result.get("error_description", result.get("error", "unknown"))
|
||||||
|
raise RuntimeError(f"MS Graph token refresh failed: {error}")
|
||||||
|
|
||||||
|
self._access_token = result["access_token"]
|
||||||
|
# MSAL may issue a new refresh token.
|
||||||
|
if "refresh_token" in result:
|
||||||
|
self._refresh_token = result["refresh_token"]
|
||||||
|
self._credentials_info["refresh_token"] = result["refresh_token"]
|
||||||
|
self._credentials_info["access_token"] = self._access_token
|
||||||
|
|
||||||
|
@property
|
||||||
|
def refreshed_credentials(self) -> dict[str, Any] | None:
|
||||||
|
"""Return updated credential dict if the access token was refreshed.
|
||||||
|
|
||||||
|
Returns ``None`` if no change was made.
|
||||||
|
"""
|
||||||
|
if self._access_token != self._original_access_token:
|
||||||
|
return {**self._credentials_info, "access_token": self._access_token}
|
||||||
|
return None
|
||||||
|
|
||||||
|
# ── HTTP helpers ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def _get(
|
||||||
|
self,
|
||||||
|
client: httpx.AsyncClient,
|
||||||
|
url: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
*,
|
||||||
|
retry_on_401: bool = True,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""GET *url* with auth; refresh token on 401 and retry once."""
|
||||||
|
resp = await client.get(url, params=params, headers=self._auth_headers())
|
||||||
|
if resp.status_code == 401 and retry_on_401 and self._refresh_token:
|
||||||
|
logger.debug("ms_graph: 401 on %s — refreshing token", url)
|
||||||
|
await self._refresh_access_token()
|
||||||
|
resp = await client.get(url, params=params, headers=self._auth_headers())
|
||||||
|
if resp.status_code == 429:
|
||||||
|
raise RuntimeError("MS Graph rate limit hit (429). Try again later.")
|
||||||
|
resp.raise_for_status()
|
||||||
|
return resp.json()
|
||||||
|
|
||||||
|
# ── Public API ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
async def fetch_emails(
|
||||||
|
self,
|
||||||
|
filter_config: dict[str, Any] | None = None,
|
||||||
|
since: datetime | None = None,
|
||||||
|
) -> list[EmailMessage]:
|
||||||
|
"""Return up to ``_MAX_EMAILS`` Outlook messages matching *filter_config*.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
filter_config:
|
||||||
|
Optional dict with ``senders``, ``date_range``, ``folders`` keys.
|
||||||
|
since:
|
||||||
|
Hard lower-bound on email date (from last agent run).
|
||||||
|
"""
|
||||||
|
odata_filter = _build_email_filter(filter_config, since)
|
||||||
|
params: dict[str, Any] = {
|
||||||
|
"$top": 50,
|
||||||
|
"$select": "id,subject,from,receivedDateTime,body,bodyPreview",
|
||||||
|
"$orderby": "receivedDateTime desc",
|
||||||
|
}
|
||||||
|
if odata_filter:
|
||||||
|
params["$filter"] = odata_filter
|
||||||
|
|
||||||
|
emails: list[EmailMessage] = []
|
||||||
|
url = f"{_GRAPH_BASE}/me/messages"
|
||||||
|
|
||||||
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||||
|
while url and len(emails) < _MAX_EMAILS:
|
||||||
|
data = await self._get(client, url, params if url.startswith(_GRAPH_BASE) else None)
|
||||||
|
for item in data.get("value", []):
|
||||||
|
emails.append(self._parse_email(item))
|
||||||
|
if len(emails) >= _MAX_EMAILS:
|
||||||
|
break
|
||||||
|
url = data.get("@odata.nextLink", "")
|
||||||
|
params = {} # nextLink already contains encoded params.
|
||||||
|
|
||||||
|
logger.info("ms_graph: fetched %d Outlook email(s)", len(emails))
|
||||||
|
return emails
|
||||||
|
|
||||||
|
async def fetch_messages(
|
||||||
|
self,
|
||||||
|
filter_config: dict[str, Any] | None = None,
|
||||||
|
since: datetime | None = None,
|
||||||
|
) -> list[ChatMessage]:
|
||||||
|
"""Return up to ``_MAX_MESSAGES`` Teams messages matching *filter_config*.
|
||||||
|
|
||||||
|
Fetches from ``/me/chats/getAllMessages`` (personal + group chats).
|
||||||
|
The ``filter_config.channels`` key is checked as a text-filter on
|
||||||
|
the channel name post-fetch (the API doesn't support channel OData
|
||||||
|
filter directly on ``getAllMessages``).
|
||||||
|
"""
|
||||||
|
cfg = filter_config or {}
|
||||||
|
channel_filter: list[str] = [c.lower() for c in cfg.get("channels", [])]
|
||||||
|
params: dict[str, Any] = {"$top": 50}
|
||||||
|
if since:
|
||||||
|
params["$filter"] = f"createdDateTime ge {_odata_datetime(since)}"
|
||||||
|
|
||||||
|
messages: list[ChatMessage] = []
|
||||||
|
url = f"{_GRAPH_BASE}/me/chats/getAllMessages"
|
||||||
|
|
||||||
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||||
|
while url and len(messages) < _MAX_MESSAGES:
|
||||||
|
try:
|
||||||
|
data = await self._get(client, url, params if url.startswith(_GRAPH_BASE) else None)
|
||||||
|
except httpx.HTTPStatusError as exc:
|
||||||
|
# getAllMessages requires specific licensing; degrade gracefully.
|
||||||
|
if exc.response.status_code in (403, 404):
|
||||||
|
logger.warning(
|
||||||
|
"ms_graph: /me/chats/getAllMessages not available (%d) — "
|
||||||
|
"check Teams license or permissions",
|
||||||
|
exc.response.status_code,
|
||||||
|
)
|
||||||
|
break
|
||||||
|
raise
|
||||||
|
|
||||||
|
for item in data.get("value", []):
|
||||||
|
msg = self._parse_teams_message(item)
|
||||||
|
if channel_filter and msg.channel:
|
||||||
|
if not any(c in msg.channel.lower() for c in channel_filter):
|
||||||
|
continue
|
||||||
|
messages.append(msg)
|
||||||
|
if len(messages) >= _MAX_MESSAGES:
|
||||||
|
break
|
||||||
|
url = data.get("@odata.nextLink", "")
|
||||||
|
params = {}
|
||||||
|
|
||||||
|
logger.info("ms_graph: fetched %d Teams message(s)", len(messages))
|
||||||
|
return messages
|
||||||
|
|
||||||
|
# ── Parsers ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _parse_email(item: dict[str, Any]) -> EmailMessage:
|
||||||
|
subject: str = item.get("subject", "(no subject)") or "(no subject)"
|
||||||
|
sender_block = item.get("from", {}) or {}
|
||||||
|
sender_addr = (
|
||||||
|
(sender_block.get("emailAddress") or {}).get("address", "unknown")
|
||||||
|
)
|
||||||
|
date_str: str = item.get("receivedDateTime", "")
|
||||||
|
try:
|
||||||
|
date = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
|
||||||
|
except Exception:
|
||||||
|
date = datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
body_block = item.get("body", {}) or {}
|
||||||
|
content_type: str = body_block.get("contentType", "text")
|
||||||
|
raw_body: str = body_block.get("content", "")
|
||||||
|
if content_type == "html":
|
||||||
|
body_text = _strip_html(raw_body)
|
||||||
|
else:
|
||||||
|
body_text = raw_body or item.get("bodyPreview", "")
|
||||||
|
body_text = body_text[:_BODY_TRUNCATE]
|
||||||
|
|
||||||
|
return EmailMessage(
|
||||||
|
id=item.get("id", ""),
|
||||||
|
subject=subject,
|
||||||
|
sender=sender_addr,
|
||||||
|
body_text=body_text,
|
||||||
|
date=date,
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _parse_teams_message(item: dict[str, Any]) -> ChatMessage:
|
||||||
|
msg_id: str = item.get("id", "")
|
||||||
|
sender_block = (item.get("from") or {}).get("user") or {}
|
||||||
|
sender: str = sender_block.get("displayName", "unknown")
|
||||||
|
channel: str | None = (item.get("channelIdentity") or {}).get("channelId")
|
||||||
|
|
||||||
|
date_str: str = item.get("createdDateTime", "")
|
||||||
|
try:
|
||||||
|
date = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
|
||||||
|
except Exception:
|
||||||
|
date = datetime.now(timezone.utc)
|
||||||
|
|
||||||
|
body_block = item.get("body", {}) or {}
|
||||||
|
content_type: str = body_block.get("contentType", "text")
|
||||||
|
raw_content: str = body_block.get("content", "")
|
||||||
|
content = _strip_html(raw_content) if content_type == "html" else raw_content
|
||||||
|
content = content[:_BODY_TRUNCATE]
|
||||||
|
|
||||||
|
return ChatMessage(
|
||||||
|
id=msg_id,
|
||||||
|
content=content,
|
||||||
|
sender=sender,
|
||||||
|
channel=channel,
|
||||||
|
date=date,
|
||||||
|
)
|
||||||
99
app/main.py
99
app/main.py
@@ -1,4 +1,5 @@
|
|||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
|
import logging
|
||||||
|
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
@@ -7,15 +8,95 @@ from app.api.middleware.rate_limit import TierRateLimitMiddleware
|
|||||||
from app.api.middleware.sanitizer import SanitizerMiddleware
|
from app.api.middleware.sanitizer import SanitizerMiddleware
|
||||||
from app.config.settings import settings
|
from app.config.settings import settings
|
||||||
|
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.INFO,
|
||||||
|
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||||||
|
)
|
||||||
|
logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)
|
||||||
|
logging.getLogger("sqlalchemy.pool").setLevel(logging.WARNING)
|
||||||
|
|
||||||
|
|
||||||
|
async def _memory_audit_cron_tick() -> None:
|
||||||
|
"""Weekly cron: contradiction scan + label canonicalization for all users (Phase 7)."""
|
||||||
|
import logging # noqa: PLC0415
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
_log.info("memory audit cron tick: starting")
|
||||||
|
try:
|
||||||
|
from app.db import async_session # noqa: PLC0415
|
||||||
|
from app.core.memory_maintenance import audit_memory # noqa: PLC0415
|
||||||
|
from app.models import User # noqa: PLC0415
|
||||||
|
from sqlalchemy import select # noqa: PLC0415
|
||||||
|
|
||||||
|
async with async_session() as db:
|
||||||
|
result = await db.execute(select(User.id))
|
||||||
|
user_ids: list[str] = list(result.scalars().all())
|
||||||
|
|
||||||
|
for uid in user_ids:
|
||||||
|
try:
|
||||||
|
async with async_session() as db:
|
||||||
|
await audit_memory(db, uid)
|
||||||
|
except Exception as exc:
|
||||||
|
_log.warning("memory audit cron tick: audit_memory failed user=%s: %s", uid, exc)
|
||||||
|
|
||||||
|
_log.info("memory audit cron tick: done users=%d", len(user_ids))
|
||||||
|
except Exception as exc:
|
||||||
|
_log.warning("memory audit cron tick: failed: %s", exc)
|
||||||
|
|
||||||
|
|
||||||
|
async def _memory_cron_tick() -> None:
|
||||||
|
"""Hourly cron: drain Free-tier extraction queue + mine proactive patterns for Power+ users."""
|
||||||
|
import logging # noqa: PLC0415
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
_log.info("memory cron tick: starting")
|
||||||
|
try:
|
||||||
|
from app.db import async_session # noqa: PLC0415
|
||||||
|
from app.core.memory_maintenance import drain_extraction_queue, mine_proactive_patterns # noqa: PLC0415
|
||||||
|
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||||
|
from app.models import User # noqa: PLC0415
|
||||||
|
from sqlalchemy import select # noqa: PLC0415
|
||||||
|
|
||||||
|
async with async_session() as db:
|
||||||
|
await drain_extraction_queue(db)
|
||||||
|
|
||||||
|
# mine proactive patterns for every Power+ user
|
||||||
|
async with async_session() as db:
|
||||||
|
result = await db.execute(select(User.id))
|
||||||
|
user_ids: list[str] = list(result.scalars().all())
|
||||||
|
|
||||||
|
for uid in user_ids:
|
||||||
|
try:
|
||||||
|
async with async_session() as db:
|
||||||
|
tier = await tier_manager.get_tier(uid, db)
|
||||||
|
if tier_manager.check_feature(tier, "proactive_mining"):
|
||||||
|
await mine_proactive_patterns(db, uid)
|
||||||
|
except Exception as exc:
|
||||||
|
_log.warning("memory cron tick: mine_proactive_patterns failed user=%s: %s", uid, exc)
|
||||||
|
|
||||||
|
_log.info("memory cron tick: done users=%d", len(user_ids))
|
||||||
|
except Exception as exc:
|
||||||
|
_log.warning("memory cron tick: failed: %s", exc)
|
||||||
|
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def lifespan(app: FastAPI):
|
||||||
# Startup: initialise DB connection pool and agent registry
|
# Startup: ensure agent tool modules are loaded.
|
||||||
from app.core.agent_registry import registry # noqa: F401 — triggers module load
|
import app.agents # noqa: F401
|
||||||
import app.agents # noqa: F401 — triggers @registry.register decorators
|
|
||||||
|
scheduler = None
|
||||||
|
if settings.SCHEDULER_ENABLED:
|
||||||
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler # noqa: PLC0415
|
||||||
|
|
||||||
|
scheduler = AsyncIOScheduler()
|
||||||
|
scheduler.add_job(_memory_cron_tick, "interval", hours=1, id="memory_cron")
|
||||||
|
scheduler.add_job(_memory_audit_cron_tick, "interval", weeks=1, id="memory_audit_cron")
|
||||||
|
scheduler.start()
|
||||||
|
logging.getLogger(__name__).info("memory cron scheduler started (interval=1h)")
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
|
if scheduler is not None:
|
||||||
|
scheduler.shutdown(wait=False)
|
||||||
|
|
||||||
# Shutdown: dispose SQLAlchemy connection pool
|
# Shutdown: dispose SQLAlchemy connection pool
|
||||||
from app.db import engine
|
from app.db import engine
|
||||||
await engine.dispose()
|
await engine.dispose()
|
||||||
@@ -23,7 +104,7 @@ async def lifespan(app: FastAPI):
|
|||||||
|
|
||||||
def create_app() -> FastAPI:
|
def create_app() -> FastAPI:
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="Adiuva Cloud API",
|
title="AdiuvAI Cloud API",
|
||||||
version="0.1.0",
|
version="0.1.0",
|
||||||
docs_url="/docs" if settings.ENV == "dev" else None,
|
docs_url="/docs" if settings.ENV == "dev" else None,
|
||||||
redoc_url=None,
|
redoc_url=None,
|
||||||
@@ -43,16 +124,14 @@ def create_app() -> FastAPI:
|
|||||||
app.add_middleware(SanitizerMiddleware)
|
app.add_middleware(SanitizerMiddleware)
|
||||||
app.add_middleware(TierRateLimitMiddleware)
|
app.add_middleware(TierRateLimitMiddleware)
|
||||||
|
|
||||||
from app.api.routes import auth, backup, billing, chat, plans, plugins, storage, vectors
|
from app.api.routes import agents, auth, billing, chat, device_ws, memory
|
||||||
|
|
||||||
app.include_router(auth.router, prefix="/api/v1")
|
app.include_router(auth.router, prefix="/api/v1")
|
||||||
app.include_router(chat.router, prefix="/api/v1")
|
app.include_router(chat.router, prefix="/api/v1")
|
||||||
app.include_router(plans.router, prefix="/api/v1")
|
|
||||||
app.include_router(storage.router, prefix="/api/v1")
|
|
||||||
app.include_router(vectors.router, prefix="/api/v1")
|
|
||||||
app.include_router(backup.router, prefix="/api/v1")
|
|
||||||
app.include_router(plugins.router, prefix="/api/v1")
|
|
||||||
app.include_router(billing.router, prefix="/api/v1")
|
app.include_router(billing.router, prefix="/api/v1")
|
||||||
|
app.include_router(agents.router, prefix="/api/v1")
|
||||||
|
app.include_router(device_ws.router, prefix="/api/v1")
|
||||||
|
app.include_router(memory.router, prefix="/api/v1")
|
||||||
|
|
||||||
@app.get("/api/v1/health", tags=["health"])
|
@app.get("/api/v1/health", tags=["health"])
|
||||||
async def health() -> dict:
|
async def health() -> dict:
|
||||||
|
|||||||
@@ -1,7 +0,0 @@
|
|||||||
"""Plugin marketplace package.
|
|
||||||
|
|
||||||
Three service classes introduced in Step 10:
|
|
||||||
- ``PluginRegistry`` — catalog, submit/approve/reject, install counts
|
|
||||||
- ``ReviewQueue`` — approval workflow + security checklist
|
|
||||||
- ``RevenueShare`` — 70/30 split tracking and Stripe Connect payouts
|
|
||||||
"""
|
|
||||||
@@ -1,212 +0,0 @@
|
|||||||
"""Plugin catalog registry backed by PostgreSQL.
|
|
||||||
|
|
||||||
Maintains the authoritative list of plugins, their review status, and
|
|
||||||
aggregate install counts. All data is persisted in the ``plugins`` table.
|
|
||||||
|
|
||||||
Module-level singleton::
|
|
||||||
|
|
||||||
from app.marketplace.plugin_registry import registry
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import json
|
|
||||||
from typing import Any, Literal
|
|
||||||
|
|
||||||
from sqlalchemy import select, func
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from app.models import Plugin
|
|
||||||
from app.schemas import PluginListResponse, PluginManifest
|
|
||||||
|
|
||||||
_PAGE_SIZE = 20
|
|
||||||
|
|
||||||
|
|
||||||
def _plugin_to_manifest(p: Plugin) -> PluginManifest:
|
|
||||||
"""Convert an ORM ``Plugin`` row to a Pydantic ``PluginManifest``."""
|
|
||||||
try:
|
|
||||||
permissions = json.loads(p.permissions) if p.permissions else []
|
|
||||||
except (json.JSONDecodeError, TypeError):
|
|
||||||
permissions = []
|
|
||||||
return PluginManifest(
|
|
||||||
id=p.id,
|
|
||||||
name=p.name,
|
|
||||||
description=p.description,
|
|
||||||
version=p.version,
|
|
||||||
author=p.author_name,
|
|
||||||
permissions=permissions,
|
|
||||||
category=p.category,
|
|
||||||
price_cents=p.price_cents,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class PluginRegistry:
|
|
||||||
"""PostgreSQL-backed plugin catalog.
|
|
||||||
|
|
||||||
All methods accept an ``AsyncSession`` parameter so the calling route
|
|
||||||
controls the session lifecycle.
|
|
||||||
"""
|
|
||||||
|
|
||||||
# ── Queries ──────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
async def list_plugins(
|
|
||||||
self,
|
|
||||||
db: AsyncSession,
|
|
||||||
category: str | None = None,
|
|
||||||
query: str | None = None,
|
|
||||||
page: int = 1,
|
|
||||||
sort: Literal["rating", "installs", "newest"] = "newest",
|
|
||||||
) -> PluginListResponse:
|
|
||||||
"""Return a page of approved plugins, optionally filtered and sorted."""
|
|
||||||
base = select(Plugin).where(Plugin.status == "approved")
|
|
||||||
|
|
||||||
if category:
|
|
||||||
base = base.where(Plugin.category == category)
|
|
||||||
if query:
|
|
||||||
pattern = f"%{query}%"
|
|
||||||
base = base.where(
|
|
||||||
Plugin.name.ilike(pattern) | Plugin.description.ilike(pattern)
|
|
||||||
)
|
|
||||||
|
|
||||||
# Count
|
|
||||||
count_q = select(func.count()).select_from(base.subquery())
|
|
||||||
total = (await db.execute(count_q)).scalar_one()
|
|
||||||
|
|
||||||
# Sort
|
|
||||||
if sort == "installs":
|
|
||||||
base = base.order_by(Plugin.install_count.desc())
|
|
||||||
elif sort == "rating":
|
|
||||||
base = base.order_by(Plugin.avg_rating.desc())
|
|
||||||
else: # newest
|
|
||||||
base = base.order_by(Plugin.created_at.desc())
|
|
||||||
|
|
||||||
base = base.offset((page - 1) * _PAGE_SIZE).limit(_PAGE_SIZE)
|
|
||||||
rows = (await db.execute(base)).scalars().all()
|
|
||||||
|
|
||||||
return PluginListResponse(
|
|
||||||
plugins=[_plugin_to_manifest(r) for r in rows],
|
|
||||||
total=total,
|
|
||||||
page=page,
|
|
||||||
)
|
|
||||||
|
|
||||||
async def get_plugin(self, db: AsyncSession, plugin_id: str) -> dict[str, Any] | None:
|
|
||||||
"""Return ``{manifest, status, install_count, avg_rating}`` or ``None``."""
|
|
||||||
result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
p = result.scalar_one_or_none()
|
|
||||||
if p is None:
|
|
||||||
return None
|
|
||||||
return {
|
|
||||||
"manifest": _plugin_to_manifest(p),
|
|
||||||
"status": p.status,
|
|
||||||
"install_count": p.install_count,
|
|
||||||
"avg_rating": p.avg_rating,
|
|
||||||
}
|
|
||||||
|
|
||||||
# ── Mutations ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
async def submit_plugin(
|
|
||||||
self,
|
|
||||||
db: AsyncSession,
|
|
||||||
manifest: PluginManifest,
|
|
||||||
package_s3_key: str,
|
|
||||||
) -> str:
|
|
||||||
"""Add *manifest* to the catalog with ``status='pending_review'``.
|
|
||||||
|
|
||||||
Returns the plugin_id. If a plugin with the same id already exists
|
|
||||||
it is overwritten (re-submission after rejection).
|
|
||||||
"""
|
|
||||||
plugin_id = manifest.id
|
|
||||||
existing = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
row = existing.scalar_one_or_none()
|
|
||||||
|
|
||||||
if row is not None:
|
|
||||||
row.name = manifest.name
|
|
||||||
row.description = manifest.description
|
|
||||||
row.version = manifest.version
|
|
||||||
row.author_name = manifest.author
|
|
||||||
row.category = manifest.category
|
|
||||||
row.price_cents = manifest.price_cents
|
|
||||||
row.permissions = json.dumps(manifest.permissions)
|
|
||||||
row.status = "pending_review"
|
|
||||||
row.s3_package_key = package_s3_key
|
|
||||||
row.rejection_reason = None
|
|
||||||
else:
|
|
||||||
row = Plugin(
|
|
||||||
id=plugin_id,
|
|
||||||
name=manifest.name,
|
|
||||||
description=manifest.description,
|
|
||||||
version=manifest.version,
|
|
||||||
author_name=manifest.author,
|
|
||||||
category=manifest.category,
|
|
||||||
price_cents=manifest.price_cents,
|
|
||||||
permissions=json.dumps(manifest.permissions),
|
|
||||||
status="pending_review",
|
|
||||||
s3_package_key=package_s3_key,
|
|
||||||
install_count=0,
|
|
||||||
avg_rating=0.0,
|
|
||||||
)
|
|
||||||
db.add(row)
|
|
||||||
await db.commit()
|
|
||||||
return plugin_id
|
|
||||||
|
|
||||||
async def approve_plugin(self, db: AsyncSession, plugin_id: str) -> None:
|
|
||||||
"""Set *plugin_id* status to ``'approved'``.
|
|
||||||
|
|
||||||
Raises ``KeyError`` if the plugin is not found.
|
|
||||||
"""
|
|
||||||
result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
row = result.scalar_one_or_none()
|
|
||||||
if row is None:
|
|
||||||
raise KeyError(f"Plugin not found: {plugin_id}")
|
|
||||||
row.status = "approved"
|
|
||||||
row.rejection_reason = None
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
async def reject_plugin(self, db: AsyncSession, plugin_id: str, reason: str) -> None:
|
|
||||||
"""Set *plugin_id* status to ``'rejected'`` and record the reason.
|
|
||||||
|
|
||||||
Raises ``KeyError`` if the plugin is not found.
|
|
||||||
"""
|
|
||||||
result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
row = result.scalar_one_or_none()
|
|
||||||
if row is None:
|
|
||||||
raise KeyError(f"Plugin not found: {plugin_id}")
|
|
||||||
row.status = "rejected"
|
|
||||||
row.rejection_reason = reason
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
async def record_install(self, db: AsyncSession, plugin_id: str) -> None:
|
|
||||||
"""Increment the install count for *plugin_id* (no-op if not found)."""
|
|
||||||
result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
row = result.scalar_one_or_none()
|
|
||||||
if row is not None:
|
|
||||||
row.install_count = row.install_count + 1
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
async def record_uninstall(self, db: AsyncSession, plugin_id: str) -> None:
|
|
||||||
"""Decrement the install count for *plugin_id*, floored at 0."""
|
|
||||||
result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
row = result.scalar_one_or_none()
|
|
||||||
if row is not None:
|
|
||||||
row.install_count = max(0, row.install_count - 1)
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
# ── Internal helpers used by ReviewQueue ─────────────────────────
|
|
||||||
|
|
||||||
async def get_pending_entries(self, db: AsyncSession) -> list[dict[str, Any]]:
|
|
||||||
"""Return all entries with status='pending_review'."""
|
|
||||||
result = await db.execute(
|
|
||||||
select(Plugin).where(Plugin.status == "pending_review")
|
|
||||||
)
|
|
||||||
rows = result.scalars().all()
|
|
||||||
return [
|
|
||||||
{
|
|
||||||
"manifest": _plugin_to_manifest(r),
|
|
||||||
"submitted_at": int(r.submitted_at.timestamp()) if r.submitted_at else 0,
|
|
||||||
}
|
|
||||||
for r in rows
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
# Module-level singleton
|
|
||||||
registry = PluginRegistry()
|
|
||||||
@@ -1,125 +0,0 @@
|
|||||||
"""Plugin review workflow backed by PostgreSQL.
|
|
||||||
|
|
||||||
Manages the approval queue for newly submitted plugins and enforces a
|
|
||||||
security checklist before any plugin is made visible in the marketplace.
|
|
||||||
|
|
||||||
Module-level singleton::
|
|
||||||
|
|
||||||
from app.marketplace.plugin_review import review_queue
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import re
|
|
||||||
from typing import Any, Literal
|
|
||||||
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from app.marketplace.plugin_registry import registry
|
|
||||||
from app.models import PluginReview as PluginReviewModel
|
|
||||||
from app.schemas import PluginManifest
|
|
||||||
|
|
||||||
# ── Security policy ───────────────────────────────────────────────────
|
|
||||||
|
|
||||||
ALLOWED_PERMISSIONS: frozenset[str] = frozenset(
|
|
||||||
{
|
|
||||||
"read:tasks",
|
|
||||||
"write:tasks",
|
|
||||||
"read:projects",
|
|
||||||
"write:projects",
|
|
||||||
"read:notes",
|
|
||||||
"write:notes",
|
|
||||||
"read:checkpoints",
|
|
||||||
"write:checkpoints",
|
|
||||||
"read:calendar",
|
|
||||||
"write:calendar",
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
_PLUGIN_ID_RE = re.compile(r"^[a-z0-9-]+$")
|
|
||||||
|
|
||||||
|
|
||||||
def validate_manifest(manifest: PluginManifest) -> None:
|
|
||||||
"""Enforce the plugin security checklist.
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
``ValueError`` on the first violation found. Callers should catch
|
|
||||||
this and return HTTP 422 / reject the submission.
|
|
||||||
|
|
||||||
Checks:
|
|
||||||
1. Plugin id matches ``^[a-z0-9-]+$``
|
|
||||||
2. All declared permissions are in ``ALLOWED_PERMISSIONS``
|
|
||||||
3. No manifest field contains raw binary data
|
|
||||||
"""
|
|
||||||
if not _PLUGIN_ID_RE.match(manifest.id):
|
|
||||||
raise ValueError(
|
|
||||||
f"Invalid plugin id format: '{manifest.id}'. "
|
|
||||||
"Only lowercase letters, digits, and hyphens are allowed."
|
|
||||||
)
|
|
||||||
|
|
||||||
for perm in manifest.permissions:
|
|
||||||
if perm not in ALLOWED_PERMISSIONS:
|
|
||||||
raise ValueError(
|
|
||||||
f"Unknown permission: '{perm}'. "
|
|
||||||
f"Allowed permissions: {sorted(ALLOWED_PERMISSIONS)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
for field_name, value in manifest.model_dump().items():
|
|
||||||
if isinstance(value, (bytes, bytearray)):
|
|
||||||
raise ValueError(
|
|
||||||
f"Binary content is not allowed in manifest field '{field_name}'."
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class ReviewQueue:
|
|
||||||
"""Approval queue for pending plugin submissions.
|
|
||||||
|
|
||||||
Delegates status changes to the shared ``PluginRegistry`` singleton.
|
|
||||||
Review records are persisted in the ``plugin_reviews`` table.
|
|
||||||
"""
|
|
||||||
|
|
||||||
async def get_pending(self, db: AsyncSession) -> list[dict[str, Any]]:
|
|
||||||
"""Return all plugins currently awaiting review.
|
|
||||||
|
|
||||||
Each item is ``{plugin_id, manifest, submitted_at}``.
|
|
||||||
"""
|
|
||||||
entries = await registry.get_pending_entries(db)
|
|
||||||
return [
|
|
||||||
{
|
|
||||||
"plugin_id": e["manifest"].id,
|
|
||||||
"manifest": e["manifest"],
|
|
||||||
"submitted_at": e["submitted_at"],
|
|
||||||
}
|
|
||||||
for e in entries
|
|
||||||
]
|
|
||||||
|
|
||||||
async def submit_review(
|
|
||||||
self,
|
|
||||||
db: AsyncSession,
|
|
||||||
plugin_id: str,
|
|
||||||
reviewer_id: str,
|
|
||||||
decision: Literal["approved", "rejected"],
|
|
||||||
notes: str = "",
|
|
||||||
) -> None:
|
|
||||||
"""Record a review decision and update the plugin's status.
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
``KeyError`` if *plugin_id* is not found in the registry.
|
|
||||||
"""
|
|
||||||
if decision == "approved":
|
|
||||||
await registry.approve_plugin(db, plugin_id)
|
|
||||||
else:
|
|
||||||
await registry.reject_plugin(db, plugin_id, reason=notes)
|
|
||||||
|
|
||||||
review = PluginReviewModel(
|
|
||||||
plugin_id=plugin_id,
|
|
||||||
reviewer_id=reviewer_id,
|
|
||||||
decision=decision,
|
|
||||||
notes=notes,
|
|
||||||
)
|
|
||||||
db.add(review)
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
|
|
||||||
# Module-level singleton
|
|
||||||
review_queue = ReviewQueue()
|
|
||||||
@@ -1,233 +0,0 @@
|
|||||||
"""Revenue share tracking and Stripe Connect payouts backed by PostgreSQL.
|
|
||||||
|
|
||||||
Records every plugin installation as a revenue event and facilitates
|
|
||||||
70 % / 30 % payouts to developers via Stripe Connect. Data is persisted
|
|
||||||
in the ``revenue_events`` table.
|
|
||||||
|
|
||||||
Module-level singleton::
|
|
||||||
|
|
||||||
from app.marketplace.revenue_share import revenue_share
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import logging
|
|
||||||
from datetime import datetime, timezone
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import stripe as stripe_lib
|
|
||||||
from sqlalchemy import extract, func, select
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from app.config.settings import settings
|
|
||||||
from app.marketplace.plugin_registry import registry
|
|
||||||
from app.models import Plugin, RevenueEvent
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
# ── Revenue split constants ───────────────────────────────────────────
|
|
||||||
|
|
||||||
DEVELOPER_SHARE: float = 0.70
|
|
||||||
PLATFORM_SHARE: float = 0.30
|
|
||||||
|
|
||||||
|
|
||||||
class RevenueShare:
|
|
||||||
"""Records installation revenue events and coordinates developer payouts.
|
|
||||||
|
|
||||||
Stripe Connect calls are gracefully stubbed when ``STRIPE_SECRET_KEY``
|
|
||||||
is not configured, consistent with the rest of the billing layer.
|
|
||||||
"""
|
|
||||||
|
|
||||||
# ── Helpers ──────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _stripe_configured() -> bool:
|
|
||||||
return bool(settings.STRIPE_SECRET_KEY)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _stripe() -> Any:
|
|
||||||
stripe_lib.api_key = settings.STRIPE_SECRET_KEY
|
|
||||||
return stripe_lib
|
|
||||||
|
|
||||||
# ── Core operations ──────────────────────────────────────────────
|
|
||||||
|
|
||||||
async def record_install(
|
|
||||||
self,
|
|
||||||
db: AsyncSession,
|
|
||||||
plugin_id: str,
|
|
||||||
user_id: str,
|
|
||||||
amount_cents: int,
|
|
||||||
) -> None:
|
|
||||||
"""Record a plugin installation and trigger a Stripe Connect charge if paid.
|
|
||||||
|
|
||||||
For free plugins (``amount_cents == 0``) no payment is initiated but
|
|
||||||
the event is still recorded for analytics.
|
|
||||||
|
|
||||||
For paid plugins the developer receives 70 % via a Stripe Connect
|
|
||||||
destination charge. If Stripe is not configured or the charge fails
|
|
||||||
the installation still succeeds (the event is recorded and the install
|
|
||||||
count is incremented) — a warning is logged for monitoring.
|
|
||||||
"""
|
|
||||||
developer_share_cents = int(amount_cents * DEVELOPER_SHARE)
|
|
||||||
stripe_transfer_id: str | None = None
|
|
||||||
|
|
||||||
if amount_cents > 0 and self._stripe_configured():
|
|
||||||
# Look up the plugin's author Stripe account from the DB
|
|
||||||
result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
plugin_row = result.scalar_one_or_none()
|
|
||||||
developer_stripe_account: str | None = None
|
|
||||||
if plugin_row and plugin_row.author_id:
|
|
||||||
# Future: look up user.stripe_connect_account_id
|
|
||||||
developer_stripe_account = None # no real account yet
|
|
||||||
|
|
||||||
if developer_stripe_account:
|
|
||||||
try:
|
|
||||||
s = self._stripe()
|
|
||||||
transfer = s.Transfer.create(
|
|
||||||
amount=developer_share_cents,
|
|
||||||
currency="eur",
|
|
||||||
destination=developer_stripe_account,
|
|
||||||
description=f"Revenue share for plugin {plugin_id}",
|
|
||||||
metadata={"plugin_id": plugin_id, "user_id": user_id},
|
|
||||||
)
|
|
||||||
stripe_transfer_id = transfer["id"]
|
|
||||||
except Exception as exc:
|
|
||||||
logger.warning(
|
|
||||||
"Stripe Connect transfer failed for plugin %s: %s",
|
|
||||||
plugin_id,
|
|
||||||
exc,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logger.debug(
|
|
||||||
"No Stripe account on file for plugin %s developer; "
|
|
||||||
"skipping transfer.",
|
|
||||||
plugin_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
event = RevenueEvent(
|
|
||||||
plugin_id=plugin_id,
|
|
||||||
user_id=user_id,
|
|
||||||
amount_cents=amount_cents,
|
|
||||||
developer_share_cents=developer_share_cents,
|
|
||||||
stripe_transfer_id=stripe_transfer_id,
|
|
||||||
)
|
|
||||||
db.add(event)
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
await registry.record_install(db, plugin_id)
|
|
||||||
|
|
||||||
async def get_earnings(
|
|
||||||
self,
|
|
||||||
db: AsyncSession,
|
|
||||||
developer_id: str,
|
|
||||||
period: str | None = None,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
"""Return aggregated earnings for *developer_id*.
|
|
||||||
|
|
||||||
``period`` is an optional ``YYYY-MM`` string to restrict the window.
|
|
||||||
|
|
||||||
Returns::
|
|
||||||
|
|
||||||
{
|
|
||||||
"developer_id": str,
|
|
||||||
"period": str | None,
|
|
||||||
"total_installs": int,
|
|
||||||
"total_revenue_cents": int,
|
|
||||||
"developer_share_cents": int,
|
|
||||||
}
|
|
||||||
"""
|
|
||||||
# Find plugin ids belonging to this developer (by author_name match)
|
|
||||||
plugin_q = select(Plugin.id).where(Plugin.author_name == developer_id)
|
|
||||||
plugin_result = await db.execute(plugin_q)
|
|
||||||
developer_plugin_ids = [row[0] for row in plugin_result.all()]
|
|
||||||
|
|
||||||
if not developer_plugin_ids:
|
|
||||||
return {
|
|
||||||
"developer_id": developer_id,
|
|
||||||
"period": period,
|
|
||||||
"total_installs": 0,
|
|
||||||
"total_revenue_cents": 0,
|
|
||||||
"developer_share_cents": 0,
|
|
||||||
}
|
|
||||||
|
|
||||||
query = select(
|
|
||||||
func.count().label("total_installs"),
|
|
||||||
func.coalesce(func.sum(RevenueEvent.amount_cents), 0).label("total_revenue"),
|
|
||||||
func.coalesce(func.sum(RevenueEvent.developer_share_cents), 0).label("dev_share"),
|
|
||||||
).where(RevenueEvent.plugin_id.in_(developer_plugin_ids))
|
|
||||||
|
|
||||||
if period:
|
|
||||||
# Filter by YYYY-MM: extract year and month from created_at
|
|
||||||
try:
|
|
||||||
year, month = period.split("-")
|
|
||||||
query = query.where(
|
|
||||||
extract("year", RevenueEvent.created_at) == int(year),
|
|
||||||
extract("month", RevenueEvent.created_at) == int(month),
|
|
||||||
)
|
|
||||||
except ValueError:
|
|
||||||
pass # invalid period format — return all
|
|
||||||
|
|
||||||
result = await db.execute(query)
|
|
||||||
row = result.one()
|
|
||||||
|
|
||||||
return {
|
|
||||||
"developer_id": developer_id,
|
|
||||||
"period": period,
|
|
||||||
"total_installs": row.total_installs,
|
|
||||||
"total_revenue_cents": row.total_revenue,
|
|
||||||
"developer_share_cents": row.dev_share,
|
|
||||||
}
|
|
||||||
|
|
||||||
async def payout_developer(self, db: AsyncSession, plugin_id: str, period: str) -> None:
|
|
||||||
"""Aggregate unpaid revenue for *period* and issue a Stripe Transfer.
|
|
||||||
|
|
||||||
Marks processed events with ``paid_at`` timestamp.
|
|
||||||
Stubs gracefully when Stripe is not configured.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
year, month = period.split("-")
|
|
||||||
year_int, month_int = int(year), int(month)
|
|
||||||
except ValueError:
|
|
||||||
logger.warning("Invalid period format: %s", period)
|
|
||||||
return
|
|
||||||
|
|
||||||
result = await db.execute(
|
|
||||||
select(RevenueEvent).where(
|
|
||||||
RevenueEvent.plugin_id == plugin_id,
|
|
||||||
RevenueEvent.paid_at.is_(None),
|
|
||||||
extract("year", RevenueEvent.created_at) == year_int,
|
|
||||||
extract("month", RevenueEvent.created_at) == month_int,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
unpaid = list(result.scalars().all())
|
|
||||||
|
|
||||||
total_dev_share = sum(e.developer_share_cents for e in unpaid)
|
|
||||||
if total_dev_share <= 0 or not unpaid:
|
|
||||||
logger.debug("Nothing to pay out for plugin %s in period %s", plugin_id, period)
|
|
||||||
return
|
|
||||||
|
|
||||||
if self._stripe_configured():
|
|
||||||
plugin_result = await db.execute(select(Plugin).where(Plugin.id == plugin_id))
|
|
||||||
plugin_row = plugin_result.scalar_one_or_none()
|
|
||||||
developer_stripe_account: str | None = None # Future: fetch from DB
|
|
||||||
if plugin_row and developer_stripe_account:
|
|
||||||
try:
|
|
||||||
s = self._stripe()
|
|
||||||
s.Transfer.create(
|
|
||||||
amount=total_dev_share,
|
|
||||||
currency="eur",
|
|
||||||
destination=developer_stripe_account,
|
|
||||||
description=f"Payout for plugin {plugin_id} period {period}",
|
|
||||||
)
|
|
||||||
except Exception as exc:
|
|
||||||
logger.warning("Payout transfer failed for plugin %s: %s", plugin_id, exc)
|
|
||||||
return
|
|
||||||
|
|
||||||
paid_ts = datetime.now(timezone.utc)
|
|
||||||
for event in unpaid:
|
|
||||||
event.paid_at = paid_ts
|
|
||||||
await db.commit()
|
|
||||||
|
|
||||||
|
|
||||||
# Module-level singleton
|
|
||||||
revenue_share = RevenueShare()
|
|
||||||
415
app/models.py
415
app/models.py
@@ -1,19 +1,20 @@
|
|||||||
"""SQLAlchemy ORM models for all persistent tables.
|
"""SQLAlchemy ORM models for all persistent tables.
|
||||||
|
|
||||||
Only auth, billing, storage metadata, and marketplace data live here.
|
Only auth, billing, agent config, and memory data live here.
|
||||||
User content (notes, tasks, etc.) is NEVER persisted server-side —
|
User content (notes, tasks, etc.) lives exclusively on the client.
|
||||||
it lives in E2E-encrypted blobs in S3, referenced by storage_records.
|
|
||||||
|
|
||||||
Table inventory:
|
Table inventory:
|
||||||
users — account credentials + tier
|
users — account credentials + tier
|
||||||
refresh_tokens — hashed refresh token store
|
refresh_tokens — hashed refresh token store
|
||||||
subscriptions — Stripe subscription records
|
subscriptions — Stripe subscription records
|
||||||
storage_records — S3 blob metadata (no plaintext)
|
local_agent_configs — per-device batch agent configs
|
||||||
backup_metadata — encrypted backup manifests
|
cloud_agent_configs — OAuth-backed cloud agent configs
|
||||||
plugins — marketplace plugin catalog
|
agent_run_logs — execution history for all agents
|
||||||
plugin_installations — per-user install records
|
memory_core — per-user persistent key/value preferences (encrypted)
|
||||||
plugin_reviews — admin review decisions
|
memory_associative — per-user semantic memory with embeddings (encrypted)
|
||||||
revenue_events — Stripe Connect 70/30 split ledger
|
memory_episodic — per-user session summaries (encrypted)
|
||||||
|
memory_proactive — per-user behavioral patterns (encrypted)
|
||||||
|
memory_relations — per-user entity/relation graph (Mem0g-light, Phase 3)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -21,16 +22,18 @@ from __future__ import annotations
|
|||||||
import uuid
|
import uuid
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
|
from pgvector.sqlalchemy import Vector
|
||||||
from sqlalchemy import (
|
from sqlalchemy import (
|
||||||
BigInteger,
|
Boolean,
|
||||||
DateTime,
|
DateTime,
|
||||||
Enum,
|
Enum,
|
||||||
Float,
|
Float,
|
||||||
ForeignKey,
|
ForeignKey,
|
||||||
Integer,
|
Integer,
|
||||||
|
JSON,
|
||||||
|
LargeBinary,
|
||||||
String,
|
String,
|
||||||
Text,
|
Text,
|
||||||
UniqueConstraint,
|
|
||||||
Uuid,
|
Uuid,
|
||||||
func,
|
func,
|
||||||
)
|
)
|
||||||
@@ -52,8 +55,9 @@ def _now() -> datetime:
|
|||||||
# ── Enum types ────────────────────────────────────────────────────────────
|
# ── Enum types ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
TierEnum = Enum("free", "pro", "power", "team", name="billing_tier")
|
TierEnum = Enum("free", "pro", "power", "team", name="billing_tier")
|
||||||
PluginStatusEnum = Enum("pending_review", "approved", "rejected", name="plugin_status")
|
AgentTypeEnum = Enum("local", "cloud", name="agent_type")
|
||||||
ReviewDecisionEnum = Enum("approved", "rejected", name="review_decision")
|
AgentStatusEnum = Enum("running", "success", "error", "partial", name="agent_run_status")
|
||||||
|
CloudProviderEnum = Enum("gmail", "teams", "outlook", name="cloud_provider")
|
||||||
|
|
||||||
|
|
||||||
# ── Models ────────────────────────────────────────────────────────────────
|
# ── Models ────────────────────────────────────────────────────────────────
|
||||||
@@ -66,12 +70,21 @@ class User(Base):
|
|||||||
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
||||||
)
|
)
|
||||||
email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True)
|
email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True)
|
||||||
password_hash: Mapped[str] = mapped_column(String(255), nullable=False)
|
name: Mapped[str | None] = mapped_column(String(100), nullable=True)
|
||||||
|
surname: Mapped[str | None] = mapped_column(String(100), nullable=True)
|
||||||
|
password_hash: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||||
|
avatar_url: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||||
tier: Mapped[str] = mapped_column(TierEnum, nullable=False, default="free")
|
tier: Mapped[str] = mapped_column(TierEnum, nullable=False, default="free")
|
||||||
stripe_customer_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
stripe_customer_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||||
|
# Per-user Fernet key (base64-urlsafe, 44 chars). Generated on registration.
|
||||||
|
# Used to encrypt/decrypt all memory rows for this user.
|
||||||
|
encryption_key: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||||
created_at: Mapped[datetime] = mapped_column(
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
)
|
)
|
||||||
|
onboarding_completed_at: Mapped[datetime | None] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=True, default=None
|
||||||
|
)
|
||||||
updated_at: Mapped[datetime] = mapped_column(
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
||||||
)
|
)
|
||||||
@@ -82,6 +95,9 @@ class User(Base):
|
|||||||
subscription: Mapped[Subscription | None] = relationship(
|
subscription: Mapped[Subscription | None] = relationship(
|
||||||
back_populates="user", uselist=False, cascade="all, delete-orphan"
|
back_populates="user", uselist=False, cascade="all, delete-orphan"
|
||||||
)
|
)
|
||||||
|
oauth_accounts: Mapped[list[OAuthAccount]] = relationship(
|
||||||
|
back_populates="user", cascade="all, delete-orphan"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class RefreshToken(Base):
|
class RefreshToken(Base):
|
||||||
@@ -102,6 +118,25 @@ class RefreshToken(Base):
|
|||||||
user: Mapped[User] = relationship(back_populates="refresh_tokens")
|
user: Mapped[User] = relationship(back_populates="refresh_tokens")
|
||||||
|
|
||||||
|
|
||||||
|
class OAuthAccount(Base):
|
||||||
|
__tablename__ = "oauth_accounts"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
||||||
|
)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
||||||
|
)
|
||||||
|
provider: Mapped[str] = mapped_column(String(50), nullable=False)
|
||||||
|
provider_user_id: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
|
provider_email: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||||
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
user: Mapped[User] = relationship(back_populates="oauth_accounts")
|
||||||
|
|
||||||
|
|
||||||
class Subscription(Base):
|
class Subscription(Base):
|
||||||
__tablename__ = "subscriptions"
|
__tablename__ = "subscriptions"
|
||||||
|
|
||||||
@@ -123,8 +158,8 @@ class Subscription(Base):
|
|||||||
user: Mapped[User] = relationship(back_populates="subscription")
|
user: Mapped[User] = relationship(back_populates="subscription")
|
||||||
|
|
||||||
|
|
||||||
class StorageRecord(Base):
|
class LocalAgentConfig(Base):
|
||||||
__tablename__ = "storage_records"
|
__tablename__ = "local_agent_configs"
|
||||||
|
|
||||||
id: Mapped[str] = mapped_column(
|
id: Mapped[str] = mapped_column(
|
||||||
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
||||||
@@ -132,10 +167,16 @@ class StorageRecord(Base):
|
|||||||
user_id: Mapped[str] = mapped_column(
|
user_id: Mapped[str] = mapped_column(
|
||||||
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
||||||
)
|
)
|
||||||
table_name: Mapped[str] = mapped_column(String(100), nullable=False)
|
device_id: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
s3_key: Mapped[str] = mapped_column(String(500), nullable=False)
|
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
checksum: Mapped[str] = mapped_column(String(64), nullable=False)
|
directory_paths: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
|
||||||
size_bytes: Mapped[int] = mapped_column(Integer, nullable=False)
|
data_types: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
|
||||||
|
prompt_template: Mapped[str] = mapped_column(Text, nullable=False, default="")
|
||||||
|
agent_config: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
||||||
|
file_extensions: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
|
||||||
|
schedule_cron: Mapped[str] = mapped_column(String(100), nullable=False, default="0 */6 * * *")
|
||||||
|
enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
|
||||||
|
last_run_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||||
created_at: Mapped[datetime] = mapped_column(
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
)
|
)
|
||||||
@@ -143,9 +184,17 @@ class StorageRecord(Base):
|
|||||||
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
||||||
)
|
)
|
||||||
|
|
||||||
|
run_logs: Mapped[list[AgentRunLog]] = relationship(
|
||||||
|
back_populates="local_agent",
|
||||||
|
primaryjoin="and_(AgentRunLog.agent_id == LocalAgentConfig.id, AgentRunLog.agent_type == 'local')",
|
||||||
|
foreign_keys="AgentRunLog.agent_id",
|
||||||
|
cascade="all, delete-orphan",
|
||||||
|
overlaps="run_logs,cloud_agent",
|
||||||
|
)
|
||||||
|
|
||||||
class BackupMetadata(Base):
|
|
||||||
__tablename__ = "backup_metadata"
|
class CloudAgentConfig(Base):
|
||||||
|
__tablename__ = "cloud_agent_configs"
|
||||||
|
|
||||||
id: Mapped[str] = mapped_column(
|
id: Mapped[str] = mapped_column(
|
||||||
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
||||||
@@ -153,116 +202,246 @@ class BackupMetadata(Base):
|
|||||||
user_id: Mapped[str] = mapped_column(
|
user_id: Mapped[str] = mapped_column(
|
||||||
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
||||||
)
|
)
|
||||||
s3_key: Mapped[str] = mapped_column(String(500), nullable=False)
|
provider: Mapped[str] = mapped_column(CloudProviderEnum, nullable=False)
|
||||||
version: Mapped[int] = mapped_column(Integer, nullable=False)
|
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
timestamp: Mapped[int] = mapped_column(BigInteger, nullable=False)
|
data_types: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
|
||||||
checksum: Mapped[str] = mapped_column(String(64), nullable=False)
|
prompt_template: Mapped[str] = mapped_column(Text, nullable=False, default="")
|
||||||
size_bytes: Mapped[int] = mapped_column(Integer, nullable=False)
|
oauth_token_encrypted: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||||
|
filter_config: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
||||||
|
schedule_cron: Mapped[str] = mapped_column(String(100), nullable=False, default="0 */6 * * *")
|
||||||
|
enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
|
||||||
|
last_run_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||||
created_at: Mapped[datetime] = mapped_column(
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
)
|
)
|
||||||
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
run_logs: Mapped[list[AgentRunLog]] = relationship(
|
||||||
|
back_populates="cloud_agent",
|
||||||
|
primaryjoin="and_(AgentRunLog.agent_id == CloudAgentConfig.id, AgentRunLog.agent_type == 'cloud')",
|
||||||
|
foreign_keys="AgentRunLog.agent_id",
|
||||||
|
cascade="all, delete-orphan",
|
||||||
|
overlaps="run_logs,local_agent",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class AgentRunLog(Base):
|
||||||
|
__tablename__ = "agent_run_logs"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
||||||
|
)
|
||||||
|
# Plain string — not a FK because it references either local_agent_configs or cloud_agent_configs
|
||||||
|
# depending on agent_type. Query by (agent_id, agent_type) to locate the source config.
|
||||||
|
agent_id: Mapped[str] = mapped_column(String(255), nullable=False, index=True)
|
||||||
|
agent_type: Mapped[str] = mapped_column(AgentTypeEnum, nullable=False)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
||||||
|
)
|
||||||
|
status: Mapped[str] = mapped_column(AgentStatusEnum, nullable=False, default="running")
|
||||||
|
items_processed: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||||
|
items_created: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||||
|
tokens_used: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
|
||||||
|
errors: Mapped[list | None] = mapped_column(JSON, nullable=True)
|
||||||
|
started_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
|
)
|
||||||
|
completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||||
|
|
||||||
|
local_agent: Mapped[LocalAgentConfig | None] = relationship(
|
||||||
|
back_populates="run_logs",
|
||||||
|
primaryjoin="and_(AgentRunLog.agent_id == LocalAgentConfig.id, AgentRunLog.agent_type == 'local')",
|
||||||
|
foreign_keys="AgentRunLog.agent_id",
|
||||||
|
overlaps="run_logs,cloud_agent",
|
||||||
|
)
|
||||||
|
cloud_agent: Mapped[CloudAgentConfig | None] = relationship(
|
||||||
|
back_populates="run_logs",
|
||||||
|
primaryjoin="and_(AgentRunLog.agent_id == CloudAgentConfig.id, AgentRunLog.agent_type == 'cloud')",
|
||||||
|
foreign_keys="AgentRunLog.agent_id",
|
||||||
|
overlaps="run_logs,local_agent",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class MonthlyTokenUsage(Base):
|
||||||
|
__tablename__ = "monthly_token_usage"
|
||||||
|
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), primary_key=True
|
||||||
|
)
|
||||||
|
year_month: Mapped[str] = mapped_column(String(7), primary_key=True) # 'YYYY-MM'
|
||||||
|
feature: Mapped[str] = mapped_column(String(64), primary_key=True)
|
||||||
|
tokens_used: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
|
||||||
|
|
||||||
|
|
||||||
|
# ── Memory models ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryCore(Base):
|
||||||
|
"""Per-user persistent key/value preferences, encrypted at rest.
|
||||||
|
|
||||||
|
Examples: preferred_language, timezone, work_style.
|
||||||
|
Decrypted in-memory only using User.encryption_key.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__tablename__ = "memory_core"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(Uuid(as_uuid=False), primary_key=True, default=_uuid)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False, index=True,
|
||||||
|
)
|
||||||
|
key: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
|
value_encrypted: Mapped[str] = mapped_column(Text, nullable=False)
|
||||||
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryAssociative(Base):
|
||||||
|
"""Per-user semantic memory: encrypted content + pgvector embedding for similarity search.
|
||||||
|
|
||||||
|
Production: ``embedding`` column is ``vector(1536)`` via pgvector.
|
||||||
|
Tests (SQLite): stored as JSON list.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__tablename__ = "memory_associative"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(Uuid(as_uuid=False), primary_key=True, default=_uuid)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False, index=True,
|
||||||
|
)
|
||||||
|
content_encrypted: Mapped[str] = mapped_column(Text, nullable=False)
|
||||||
|
# vector(1536) via pgvector; SQLite tests use NULL embeddings so no dialect issue.
|
||||||
|
embedding: Mapped[list | None] = mapped_column(Vector(1536), nullable=True)
|
||||||
|
entity_type: Mapped[str | None] = mapped_column(String(100), nullable=True)
|
||||||
|
entity_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||||
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryEpisodic(Base):
|
||||||
|
"""Per-user session summaries, encrypted at rest.
|
||||||
|
|
||||||
|
One row per session interaction; used to recall recent conversations.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__tablename__ = "memory_episodic"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(Uuid(as_uuid=False), primary_key=True, default=_uuid)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False, index=True,
|
||||||
|
)
|
||||||
|
summary_encrypted: Mapped[str] = mapped_column(Text, nullable=False)
|
||||||
|
session_id: Mapped[str] = mapped_column(String(255), nullable=False, index=True)
|
||||||
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryProactive(Base):
|
||||||
|
"""Per-user inferred behavioral patterns, encrypted at rest.
|
||||||
|
|
||||||
|
Confidence in [0.0, 1.0]; only patterns above threshold are injected.
|
||||||
|
Source: 'inferred' (from episodes) or 'explicit' (user-stated).
|
||||||
|
"""
|
||||||
|
|
||||||
|
__tablename__ = "memory_proactive"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(Uuid(as_uuid=False), primary_key=True, default=_uuid)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False, index=True,
|
||||||
|
)
|
||||||
|
pattern_encrypted: Mapped[str] = mapped_column(Text, nullable=False)
|
||||||
|
confidence: Mapped[float] = mapped_column(Float, nullable=False, default=0.5)
|
||||||
|
source: Mapped[str] = mapped_column(String(50), nullable=False, default="inferred")
|
||||||
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class ExtractionQueue(Base):
|
||||||
|
"""Batch extraction queue for Free-tier users (Phase 2).
|
||||||
|
|
||||||
|
Pro/Power/Team users get realtime asyncio.create_task() extraction.
|
||||||
|
Free users get a queue row here; a daily cron (Phase 5) drains it.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__tablename__ = "extraction_queue"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(Uuid(as_uuid=False), primary_key=True, default=_uuid)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False, index=True,
|
||||||
|
)
|
||||||
|
episode_id: Mapped[str | None] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), nullable=True,
|
||||||
|
)
|
||||||
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryRelation(Base):
|
||||||
|
"""Per-user entity/relation graph row (Mem0g-light, Phase 3).
|
||||||
|
|
||||||
|
subject_label/object_label are plaintext entity identifiers (not user content).
|
||||||
|
notes_encrypted is optional Fernet-encrypted per-user commentary.
|
||||||
|
confidence in [0.0, 1.0] — decays 5 % per 30 days since last_confirmed_at.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__tablename__ = "memory_relations"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(Uuid(as_uuid=False), primary_key=True, default=_uuid)
|
||||||
|
user_id: Mapped[str] = mapped_column(
|
||||||
|
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"),
|
||||||
|
nullable=False, index=True,
|
||||||
|
)
|
||||||
|
subject_label: Mapped[str] = mapped_column(String(128), nullable=False)
|
||||||
|
subject_type: Mapped[str] = mapped_column(String(32), nullable=False)
|
||||||
|
predicate: Mapped[str] = mapped_column(String(64), nullable=False)
|
||||||
|
object_label: Mapped[str] = mapped_column(String(128), nullable=False)
|
||||||
|
object_type: Mapped[str] = mapped_column(String(32), nullable=False)
|
||||||
|
confidence: Mapped[float] = mapped_column(Float, nullable=False, default=0.7)
|
||||||
|
source_episode_id: Mapped[str | None] = mapped_column(
|
||||||
|
Uuid(as_uuid=False),
|
||||||
|
ForeignKey("memory_episodic.id", ondelete="SET NULL"),
|
||||||
|
nullable=True,
|
||||||
|
)
|
||||||
|
notes_encrypted: Mapped[bytes | None] = mapped_column(LargeBinary, nullable=True)
|
||||||
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
|
)
|
||||||
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
|
||||||
|
)
|
||||||
|
last_confirmed_at: Mapped[datetime | None] = mapped_column(
|
||||||
|
DateTime(timezone=True), nullable=True
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class Plugin(Base):
|
class Plugin(Base):
|
||||||
|
"""Plugin marketplace catalog entry."""
|
||||||
|
|
||||||
__tablename__ = "plugins"
|
__tablename__ = "plugins"
|
||||||
|
|
||||||
id: Mapped[str] = mapped_column(String(255), primary_key=True)
|
id: Mapped[str] = mapped_column(String(255), primary_key=True)
|
||||||
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
description: Mapped[str] = mapped_column(Text, nullable=False, default="")
|
description: Mapped[str] = mapped_column(Text, nullable=False)
|
||||||
version: Mapped[str] = mapped_column(String(50), nullable=False, default="1.0.0")
|
version: Mapped[str] = mapped_column(String(50), nullable=False)
|
||||||
# nullable until developer account system is built
|
author_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||||
author_id: Mapped[str | None] = mapped_column(
|
category: Mapped[str] = mapped_column(String(100), nullable=False)
|
||||||
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
|
||||||
)
|
|
||||||
author_name: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
|
||||||
category: Mapped[str] = mapped_column(String(100), nullable=False, default="")
|
|
||||||
price_cents: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
price_cents: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||||
permissions: Mapped[str] = mapped_column(Text, nullable=False, default="[]") # JSON list
|
permissions: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
|
||||||
status: Mapped[str] = mapped_column(PluginStatusEnum, nullable=False, default="pending_review")
|
status: Mapped[str] = mapped_column(String(50), nullable=False, default="pending")
|
||||||
s3_package_key: Mapped[str | None] = mapped_column(String(500), nullable=True)
|
s3_package_key: Mapped[str | None] = mapped_column(String(500), nullable=True)
|
||||||
install_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
install_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||||
avg_rating: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
|
avg_rating: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
|
||||||
rejection_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
||||||
submitted_at: Mapped[datetime] = mapped_column(
|
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
|
||||||
)
|
|
||||||
created_at: Mapped[datetime] = mapped_column(
|
created_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||||
)
|
)
|
||||||
|
|
||||||
installations: Mapped[list[PluginInstallation]] = relationship(
|
|
||||||
back_populates="plugin", cascade="all, delete-orphan"
|
|
||||||
)
|
|
||||||
reviews: Mapped[list[PluginReview]] = relationship(
|
|
||||||
back_populates="plugin", cascade="all, delete-orphan"
|
|
||||||
)
|
|
||||||
revenue_events: Mapped[list[RevenueEvent]] = relationship(
|
|
||||||
back_populates="plugin", cascade="all, delete-orphan"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class PluginInstallation(Base):
|
|
||||||
__tablename__ = "plugin_installations"
|
|
||||||
__table_args__ = (UniqueConstraint("plugin_id", "user_id", name="uq_plugin_user"),)
|
|
||||||
|
|
||||||
id: Mapped[str] = mapped_column(
|
|
||||||
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
|
||||||
)
|
|
||||||
plugin_id: Mapped[str] = mapped_column(
|
|
||||||
String(255), ForeignKey("plugins.id", ondelete="CASCADE"), nullable=False, index=True
|
|
||||||
)
|
|
||||||
user_id: Mapped[str] = mapped_column(
|
|
||||||
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
|
||||||
)
|
|
||||||
installed_at: Mapped[datetime] = mapped_column(
|
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
|
||||||
)
|
|
||||||
|
|
||||||
plugin: Mapped[Plugin] = relationship(back_populates="installations")
|
|
||||||
|
|
||||||
|
|
||||||
class PluginReview(Base):
|
|
||||||
__tablename__ = "plugin_reviews"
|
|
||||||
|
|
||||||
id: Mapped[str] = mapped_column(
|
|
||||||
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
|
||||||
)
|
|
||||||
plugin_id: Mapped[str] = mapped_column(
|
|
||||||
String(255), ForeignKey("plugins.id", ondelete="CASCADE"), nullable=False, index=True
|
|
||||||
)
|
|
||||||
reviewer_id: Mapped[str | None] = mapped_column(
|
|
||||||
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
|
||||||
)
|
|
||||||
decision: Mapped[str] = mapped_column(ReviewDecisionEnum, nullable=False)
|
|
||||||
notes: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
||||||
reviewed_at: Mapped[datetime] = mapped_column(
|
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
|
||||||
)
|
|
||||||
created_at: Mapped[datetime] = mapped_column(
|
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
|
||||||
)
|
|
||||||
|
|
||||||
plugin: Mapped[Plugin] = relationship(back_populates="reviews")
|
|
||||||
|
|
||||||
|
|
||||||
class RevenueEvent(Base):
|
|
||||||
__tablename__ = "revenue_events"
|
|
||||||
|
|
||||||
id: Mapped[str] = mapped_column(
|
|
||||||
Uuid(as_uuid=False), primary_key=True, default=_uuid
|
|
||||||
)
|
|
||||||
plugin_id: Mapped[str] = mapped_column(
|
|
||||||
String(255), ForeignKey("plugins.id", ondelete="CASCADE"), nullable=False, index=True
|
|
||||||
)
|
|
||||||
user_id: Mapped[str] = mapped_column(
|
|
||||||
Uuid(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
|
|
||||||
)
|
|
||||||
amount_cents: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
||||||
developer_share_cents: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
||||||
stripe_transfer_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
|
||||||
paid_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
||||||
created_at: Mapped[datetime] = mapped_column(
|
|
||||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
|
||||||
)
|
|
||||||
|
|
||||||
plugin: Mapped[Plugin] = relationship(back_populates="revenue_events")
|
|
||||||
|
|||||||
318
app/schemas.py
318
app/schemas.py
@@ -27,7 +27,19 @@ class AuthTokens(BaseModel):
|
|||||||
class UserProfile(BaseModel):
|
class UserProfile(BaseModel):
|
||||||
id: str
|
id: str
|
||||||
email: str
|
email: str
|
||||||
|
name: str | None = None
|
||||||
|
surname: str | None = None
|
||||||
tier: BillingTier
|
tier: BillingTier
|
||||||
|
avatar_url: str | None = None
|
||||||
|
has_password: bool = True
|
||||||
|
onboarding_completed_at: int | None = None # epoch ms, null = not onboarded
|
||||||
|
memory: dict[str, str] = Field(default_factory=dict) # decrypted core memory k/v
|
||||||
|
|
||||||
|
|
||||||
|
class OAuthAccountInfo(BaseModel):
|
||||||
|
provider: str
|
||||||
|
provider_email: str | None = None
|
||||||
|
created_at: int # epoch ms
|
||||||
|
|
||||||
|
|
||||||
# ── Chat ─────────────────────────────────────────────────────────────
|
# ── Chat ─────────────────────────────────────────────────────────────
|
||||||
@@ -39,134 +51,51 @@ class ChatContext(BaseModel):
|
|||||||
conversation_history: list[dict[str, Any]] = Field(default_factory=list)
|
conversation_history: list[dict[str, Any]] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
class PlanAction(BaseModel):
|
|
||||||
type: Literal[
|
|
||||||
"create_record",
|
|
||||||
"update_record",
|
|
||||||
"delete_record",
|
|
||||||
"index_document",
|
|
||||||
"send_notification",
|
|
||||||
]
|
|
||||||
table: str | None = None
|
|
||||||
data: dict[str, Any] | None = None
|
|
||||||
|
|
||||||
|
|
||||||
class ChatRequest(BaseModel):
|
class ChatRequest(BaseModel):
|
||||||
message: str
|
message: str
|
||||||
context: ChatContext = Field(default_factory=ChatContext)
|
context: ChatContext = Field(default_factory=ChatContext)
|
||||||
execution_mode: Literal["direct", "plan"] = "direct"
|
|
||||||
|
|
||||||
|
|
||||||
class ChatResponse(BaseModel):
|
class ChatResponse(BaseModel):
|
||||||
response: str
|
response: str
|
||||||
actions: list[PlanAction] = Field(default_factory=list)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Execution Plans ──────────────────────────────────────────────────
|
|
||||||
|
|
||||||
class PlanStep(BaseModel):
|
|
||||||
action: str
|
|
||||||
prompt_template: str | None = None
|
|
||||||
variables: dict[str, Any] | None = None
|
|
||||||
data_from_step: int | None = None
|
|
||||||
|
|
||||||
|
|
||||||
class ExecutionPlan(BaseModel):
|
|
||||||
agent: str
|
|
||||||
steps: list[PlanStep] = Field(default_factory=list)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Backup ───────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
class BackupMetadata(BaseModel):
|
|
||||||
version: int
|
|
||||||
timestamp: int
|
|
||||||
checksum: str
|
|
||||||
chunk_count: int
|
|
||||||
|
|
||||||
|
|
||||||
# ── Cloud Storage (E2E encrypted blobs) ──────────────────────────────
|
|
||||||
|
|
||||||
class StorageRecord(BaseModel):
|
|
||||||
id: str
|
|
||||||
user_id: str
|
|
||||||
table: str
|
|
||||||
blob: bytes
|
|
||||||
checksum: str
|
|
||||||
created_at: int
|
|
||||||
updated_at: int
|
|
||||||
|
|
||||||
|
|
||||||
class StorageRecordCreate(BaseModel):
|
|
||||||
table: str
|
|
||||||
blob: bytes
|
|
||||||
checksum: str
|
|
||||||
|
|
||||||
|
|
||||||
class StorageRecordUpdate(BaseModel):
|
|
||||||
blob: bytes
|
|
||||||
checksum: str
|
|
||||||
|
|
||||||
|
|
||||||
# ── Cloud Vector Store (E2E encrypted vectors) ────────────────────────
|
|
||||||
|
|
||||||
class VectorItem(BaseModel):
|
|
||||||
id: str
|
|
||||||
blob: bytes # encrypted vector + metadata — backend never decrypts
|
|
||||||
checksum: str
|
|
||||||
|
|
||||||
|
|
||||||
class VectorUpsertRequest(BaseModel):
|
|
||||||
vectors: list[VectorItem]
|
|
||||||
|
|
||||||
|
|
||||||
class VectorSearchRequest(BaseModel):
|
|
||||||
query_blob: bytes # encrypted query — backend never decrypts
|
|
||||||
top_k: int = 10
|
|
||||||
|
|
||||||
|
|
||||||
class VectorSearchResult(BaseModel):
|
|
||||||
id: str
|
|
||||||
score: float
|
|
||||||
blob: bytes
|
|
||||||
|
|
||||||
|
|
||||||
class VectorSearchResponse(BaseModel):
|
|
||||||
results: list[VectorSearchResult]
|
|
||||||
|
|
||||||
|
|
||||||
# ── Plugin Marketplace ────────────────────────────────────────────────
|
|
||||||
|
|
||||||
class PluginManifest(BaseModel):
|
|
||||||
id: str
|
|
||||||
name: str
|
|
||||||
description: str
|
|
||||||
version: str
|
|
||||||
author: str
|
|
||||||
permissions: list[str]
|
|
||||||
category: str
|
|
||||||
price_cents: int = 0
|
|
||||||
|
|
||||||
|
|
||||||
class PluginListResponse(BaseModel):
|
|
||||||
plugins: list[PluginManifest]
|
|
||||||
total: int
|
|
||||||
page: int
|
|
||||||
|
|
||||||
|
|
||||||
class PluginInstallRequest(BaseModel):
|
|
||||||
plugin_id: str
|
|
||||||
|
|
||||||
|
|
||||||
# ── WebSocket Frame Protocol ──────────────────────────────────────────
|
# ── WebSocket Frame Protocol ──────────────────────────────────────────
|
||||||
|
|
||||||
class WsFrameType(str, Enum):
|
class WsFrameType(str, Enum):
|
||||||
|
# ── v2 frame types (kept for backward compat) ──────────────────────
|
||||||
chat_request = "chat_request"
|
chat_request = "chat_request"
|
||||||
text_chunk = "text_chunk"
|
text_chunk = "text_chunk"
|
||||||
tool_call = "tool_call"
|
tool_call = "tool_call"
|
||||||
tool_result = "tool_result"
|
tool_result = "tool_result"
|
||||||
final = "final"
|
final = "final"
|
||||||
ping = "ping"
|
ping = "ping"
|
||||||
|
device_hello = "device_hello"
|
||||||
|
# ── v3 frame types ─────────────────────────────────────────────────
|
||||||
|
home_request = "home_request"
|
||||||
|
floating_request = "floating_request"
|
||||||
|
stream_start = "stream_start"
|
||||||
|
stream_text = "stream_text"
|
||||||
|
stream_end = "stream_end"
|
||||||
|
floating_domain = "floating_domain"
|
||||||
|
data_request = "data_request"
|
||||||
|
data_response = "data_response"
|
||||||
|
mutation = "mutation"
|
||||||
|
# ── v4 journey frame types ────────────────────────────────────────
|
||||||
|
journey_start = "journey_start"
|
||||||
|
journey_message = "journey_message"
|
||||||
|
journey_reply = "journey_reply"
|
||||||
|
# ── v5 brief frame types ──────────────────────────────────────────
|
||||||
|
brief_request = "brief_request"
|
||||||
|
# ── v6 task brief frame types ─────────────────────────────────────
|
||||||
|
task_brief_request = "task_brief_request"
|
||||||
|
# ── v7 folder index frame types ───────────────────────────────────
|
||||||
|
index_session_start = "index_session_start"
|
||||||
|
index_file_batch = "index_file_batch"
|
||||||
|
index_session_cancel = "index_session_cancel"
|
||||||
|
index_file_result = "index_file_result"
|
||||||
|
index_session_progress = "index_session_progress"
|
||||||
|
index_session_done = "index_session_done"
|
||||||
|
|
||||||
|
|
||||||
class WsToolCall(BaseModel):
|
class WsToolCall(BaseModel):
|
||||||
@@ -207,3 +136,172 @@ class WsFinal(BaseModel):
|
|||||||
|
|
||||||
type: Literal[WsFrameType.final] = WsFrameType.final
|
type: Literal[WsFrameType.final] = WsFrameType.final
|
||||||
response: str
|
response: str
|
||||||
|
|
||||||
|
|
||||||
|
# ── WebSocket Agent Frame Protocol ────────────────────────────────────
|
||||||
|
|
||||||
|
class WsDeviceHello(BaseModel):
|
||||||
|
"""Client → Server: device identification on WS connect."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.device_hello] = WsFrameType.device_hello
|
||||||
|
device_id: str
|
||||||
|
agent_ids: list[str] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
# ── WebSocket v3 Frame Models ─────────────────────────────────────────
|
||||||
|
|
||||||
|
class FormatPrefsModel(BaseModel):
|
||||||
|
"""User display preferences sent by Electron on each request."""
|
||||||
|
|
||||||
|
timezone: str = "UTC"
|
||||||
|
date_format: str = "dd/MM/yyyy"
|
||||||
|
time_format: str = "24h"
|
||||||
|
locale: str = "en-US"
|
||||||
|
now_iso: str = ""
|
||||||
|
|
||||||
|
|
||||||
|
class WsFloatingScope(BaseModel):
|
||||||
|
"""Scope for a floating request — narrows the agent to a specific entity."""
|
||||||
|
|
||||||
|
type: Literal["task", "project", "note", "timeline"]
|
||||||
|
id: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WsHomeRequest(BaseModel):
|
||||||
|
"""Client → Server: Home chat message."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.home_request] = WsFrameType.home_request
|
||||||
|
message: str
|
||||||
|
conversation_history: list[dict[str, Any]] = Field(default_factory=list)
|
||||||
|
format_prefs: FormatPrefsModel | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WsFloatingRequest(BaseModel):
|
||||||
|
"""Client → Server: Floating chat message scoped to an entity."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.floating_request] = WsFrameType.floating_request
|
||||||
|
message: str
|
||||||
|
scope: WsFloatingScope
|
||||||
|
format_prefs: FormatPrefsModel | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WsBriefRequest(BaseModel):
|
||||||
|
"""Client → Server: Request a plain-text brief (home or project)."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.brief_request] = WsFrameType.brief_request
|
||||||
|
request_id: str | None = None
|
||||||
|
session_id: str | None = None
|
||||||
|
mode: Literal["home", "project"]
|
||||||
|
project_id: str | None = None
|
||||||
|
format_prefs: FormatPrefsModel | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WsStreamStart(BaseModel):
|
||||||
|
"""Server → Client: signals start of a streaming response."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.stream_start] = WsFrameType.stream_start
|
||||||
|
request_id: str
|
||||||
|
|
||||||
|
|
||||||
|
class WsStreamText(BaseModel):
|
||||||
|
"""Server → Client: streamed text token."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.stream_text] = WsFrameType.stream_text
|
||||||
|
request_id: str
|
||||||
|
chunk: str
|
||||||
|
|
||||||
|
|
||||||
|
class WsStreamEnd(BaseModel):
|
||||||
|
"""Server → Client: signals end of a streaming response."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.stream_end] = WsFrameType.stream_end
|
||||||
|
request_id: str
|
||||||
|
error: str | None = None
|
||||||
|
mutations: list[dict[str, Any]] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WsDomain(BaseModel):
|
||||||
|
"""Structured floating domain payload for UI routing decisions."""
|
||||||
|
|
||||||
|
type: Literal["task", "timeline", "project", "node"]
|
||||||
|
id: str | None = None
|
||||||
|
section: Literal["task", "timeline", "note"] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WsFloatingDomain(BaseModel):
|
||||||
|
"""Server → Client: domain determined for a floating request."""
|
||||||
|
|
||||||
|
type: Literal[WsFrameType.floating_domain] = WsFrameType.floating_domain
|
||||||
|
request_id: str
|
||||||
|
domain: WsDomain
|
||||||
|
|
||||||
|
|
||||||
|
# ── Agent Config V2 ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class ContentTypeConfig(BaseModel):
|
||||||
|
"""Per-type extraction config produced by the journey chatbot."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
label: str = ""
|
||||||
|
detection_hint: str = ""
|
||||||
|
preprocessing: str = "generic" # handler name: "email_html", "plain_text", ...
|
||||||
|
extraction_prompt: str
|
||||||
|
|
||||||
|
|
||||||
|
class AgentConfig(BaseModel):
|
||||||
|
"""Structured agent configuration (replaces freeform prompt_template)."""
|
||||||
|
|
||||||
|
content_types: list[ContentTypeConfig] = []
|
||||||
|
global_rules: list[str] = []
|
||||||
|
data_types: list[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
# ── Agent Catalog ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
class AgentCatalogItem(BaseModel):
|
||||||
|
type: str
|
||||||
|
name: str
|
||||||
|
description: str
|
||||||
|
|
||||||
|
|
||||||
|
class AgentCreationCheckRequest(BaseModel):
|
||||||
|
active_agents: int = Field(ge=0, default=0)
|
||||||
|
|
||||||
|
|
||||||
|
class AgentCreationCheckResponse(BaseModel):
|
||||||
|
allowed: bool
|
||||||
|
tier: BillingTier
|
||||||
|
active_agents: int
|
||||||
|
limit: int
|
||||||
|
|
||||||
|
|
||||||
|
class AgentTriggerRequest(BaseModel):
|
||||||
|
directory: str = Field(min_length=1)
|
||||||
|
device_id: str = Field(default="")
|
||||||
|
agent_id: str | None = None # FE stable agent ID (electron-store UUID)
|
||||||
|
what_to_extract: list[str] = Field(min_length=1)
|
||||||
|
batch_interval: str = Field(min_length=1)
|
||||||
|
custom_agent_prompt: str | None = None
|
||||||
|
agent_config: dict | None = None
|
||||||
|
active_agents: int = Field(ge=0, default=0)
|
||||||
|
last_run_at: int | None = None # epoch ms from FE — enables incremental scanning
|
||||||
|
|
||||||
|
|
||||||
|
# ── Agent Run Log ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
class AgentRunLogResponse(BaseModel):
|
||||||
|
id: str
|
||||||
|
agent_id: str
|
||||||
|
agent_type: Literal["local", "cloud"]
|
||||||
|
status: Literal["running", "success", "error", "partial"]
|
||||||
|
items_processed: int
|
||||||
|
items_created: int
|
||||||
|
errors: list[str]
|
||||||
|
started_at: int
|
||||||
|
completed_at: int | None
|
||||||
|
|
||||||
|
|
||||||
|
# ── Chatbot Journey ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|||||||
@@ -1 +0,0 @@
|
|||||||
"""Cloud storage layer — E2E encrypted blobs and vectors."""
|
|
||||||
@@ -1,106 +0,0 @@
|
|||||||
"""S3-backed store for E2E-encrypted blobs.
|
|
||||||
|
|
||||||
Keys are structured as ``{user_id}/{table}/{record_id}``.
|
|
||||||
The backend never inspects blob content — it stores and retrieves opaque bytes.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import boto3
|
|
||||||
|
|
||||||
from app.config.settings import settings
|
|
||||||
|
|
||||||
|
|
||||||
class BlobStore:
|
|
||||||
"""Thin wrapper around boto3 S3.
|
|
||||||
|
|
||||||
All blobs must be E2E encrypted by the client before upload.
|
|
||||||
The backend adds SSE-S3 as an extra layer of at-rest encryption
|
|
||||||
but cannot decrypt the inner client-side payload.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def _client(self) -> Any:
|
|
||||||
kwargs: dict[str, Any] = {
|
|
||||||
"region_name": settings.S3_REGION,
|
|
||||||
"aws_access_key_id": settings.AWS_ACCESS_KEY_ID,
|
|
||||||
"aws_secret_access_key": settings.AWS_SECRET_ACCESS_KEY,
|
|
||||||
}
|
|
||||||
if settings.S3_ENDPOINT_URL and isinstance(settings.S3_ENDPOINT_URL, str):
|
|
||||||
kwargs["endpoint_url"] = settings.S3_ENDPOINT_URL
|
|
||||||
return boto3.client("s3", **kwargs)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _key(user_id: str, table: str, record_id: str) -> str:
|
|
||||||
return f"{user_id}/{table}/{record_id}"
|
|
||||||
|
|
||||||
async def upload(
|
|
||||||
self,
|
|
||||||
user_id: str,
|
|
||||||
table: str,
|
|
||||||
record_id: str,
|
|
||||||
blob: bytes,
|
|
||||||
checksum: str,
|
|
||||||
) -> str:
|
|
||||||
"""Store *blob* in S3 and return the S3 key.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
user_id: Owner of the blob (used as key prefix).
|
|
||||||
table: Logical table name (e.g. ``"tasks"``).
|
|
||||||
record_id: Record UUID.
|
|
||||||
blob: Raw bytes (pre-encrypted by client).
|
|
||||||
checksum: SHA-256 hex digest supplied by the client; stored as
|
|
||||||
object metadata for download-time verification.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
The S3 key under which the blob was stored.
|
|
||||||
"""
|
|
||||||
key = self._key(user_id, table, record_id)
|
|
||||||
self._client().put_object(
|
|
||||||
Bucket=settings.S3_BUCKET,
|
|
||||||
Key=key,
|
|
||||||
Body=blob,
|
|
||||||
ServerSideEncryption="AES256", # SSE-S3 at rest
|
|
||||||
Metadata={"checksum": checksum},
|
|
||||||
)
|
|
||||||
return key
|
|
||||||
|
|
||||||
async def download(self, user_id: str, s3_key: str) -> bytes:
|
|
||||||
"""Retrieve the blob stored at *s3_key*.
|
|
||||||
|
|
||||||
*user_id* is retained in the signature so higher-level code can
|
|
||||||
enforce ownership without re-parsing the key.
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
``botocore.exceptions.ClientError`` with code ``NoSuchKey`` if the
|
|
||||||
object does not exist.
|
|
||||||
"""
|
|
||||||
response = self._client().get_object(
|
|
||||||
Bucket=settings.S3_BUCKET,
|
|
||||||
Key=s3_key,
|
|
||||||
)
|
|
||||||
return response["Body"].read()
|
|
||||||
|
|
||||||
async def delete(self, user_id: str, s3_key: str) -> None:
|
|
||||||
"""Delete the object at *s3_key*.
|
|
||||||
|
|
||||||
S3 ``delete_object`` is idempotent — it succeeds even if the key does
|
|
||||||
not exist.
|
|
||||||
"""
|
|
||||||
self._client().delete_object(
|
|
||||||
Bucket=settings.S3_BUCKET,
|
|
||||||
Key=s3_key,
|
|
||||||
)
|
|
||||||
|
|
||||||
async def list_keys(self, user_id: str, table: str) -> list[str]:
|
|
||||||
"""Return all S3 keys for a given user + table combination.
|
|
||||||
|
|
||||||
Uses the prefix ``{user_id}/{table}/`` to scope the listing.
|
|
||||||
"""
|
|
||||||
prefix = f"{user_id}/{table}/"
|
|
||||||
response = self._client().list_objects_v2(
|
|
||||||
Bucket=settings.S3_BUCKET,
|
|
||||||
Prefix=prefix,
|
|
||||||
)
|
|
||||||
return [obj["Key"] for obj in response.get("Contents", [])]
|
|
||||||
@@ -1,32 +0,0 @@
|
|||||||
"""Integrity verification only — the backend NEVER decrypts user data."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import hashlib
|
|
||||||
import hmac
|
|
||||||
|
|
||||||
from fastapi import HTTPException
|
|
||||||
|
|
||||||
|
|
||||||
def verify_checksum(blob: bytes, checksum: str) -> bool:
|
|
||||||
"""Return ``True`` if SHA-256(blob) matches *checksum*.
|
|
||||||
|
|
||||||
Uses ``hmac.compare_digest`` for constant-time comparison to prevent
|
|
||||||
timing-based side-channel attacks.
|
|
||||||
"""
|
|
||||||
computed = hashlib.sha256(blob).hexdigest()
|
|
||||||
return hmac.compare_digest(computed, checksum)
|
|
||||||
|
|
||||||
|
|
||||||
def reject_if_tampered(blob: bytes, checksum: str) -> None:
|
|
||||||
"""Raise ``HTTP 400`` if the blob does not match its checksum.
|
|
||||||
|
|
||||||
Call this before storing or forwarding any client-provided blob.
|
|
||||||
The backend never holds decryption keys — this check only verifies
|
|
||||||
that the opaque bytes arrived intact.
|
|
||||||
"""
|
|
||||||
if not verify_checksum(blob, checksum):
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=400,
|
|
||||||
detail="Checksum mismatch: blob integrity check failed",
|
|
||||||
)
|
|
||||||
@@ -1,205 +0,0 @@
|
|||||||
"""Cloud vector store — wraps Pinecone (default) or Qdrant.
|
|
||||||
|
|
||||||
Vectors are pre-encrypted blobs from the client. The backend stores them
|
|
||||||
alongside a deterministic 32-dim float representation derived from the blob's
|
|
||||||
SHA-256 hash. Semantic ANN search is not meaningful on encrypted data — this
|
|
||||||
is a known trade-off documented in the backend plan.
|
|
||||||
|
|
||||||
Isolation: Pinecone uses ``namespace=user_id``; Qdrant filters by
|
|
||||||
``user_id`` payload field on a shared collection.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import base64
|
|
||||||
import hashlib
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
from pinecone import Pinecone
|
|
||||||
from qdrant_client import QdrantClient
|
|
||||||
from qdrant_client.models import FieldCondition, Filter, MatchValue, PointIdsList, PointStruct
|
|
||||||
|
|
||||||
from app.config.settings import settings
|
|
||||||
from app.schemas import VectorItem, VectorSearchResult
|
|
||||||
|
|
||||||
_QDRANT_COLLECTION = "adiuva_vectors"
|
|
||||||
|
|
||||||
|
|
||||||
def _blob_to_vector(blob: bytes) -> list[float]:
|
|
||||||
"""Derive a 32-dim float vector from *blob* for storage purposes only.
|
|
||||||
|
|
||||||
Uses SHA-256 to produce a deterministic 32-byte fingerprint, then
|
|
||||||
normalises each byte to the range [-1.0, 1.0]. This vector carries no
|
|
||||||
semantic meaning on encrypted data.
|
|
||||||
"""
|
|
||||||
return [(b - 128) / 128.0 for b in hashlib.sha256(blob).digest()]
|
|
||||||
|
|
||||||
|
|
||||||
class VectorStore:
|
|
||||||
"""Thin wrapper around Pinecone or Qdrant.
|
|
||||||
|
|
||||||
The backend to use is selected at runtime:
|
|
||||||
- Pinecone: when ``settings.PINECONE_API_KEY`` is non-empty.
|
|
||||||
- Qdrant: otherwise (requires ``settings.QDRANT_URL``).
|
|
||||||
"""
|
|
||||||
|
|
||||||
def _use_pinecone(self) -> bool:
|
|
||||||
return bool(settings.PINECONE_API_KEY)
|
|
||||||
|
|
||||||
# ── Pinecone helpers ──────────────────────────────────────────────
|
|
||||||
|
|
||||||
def _pinecone_index(self) -> Any:
|
|
||||||
pc = Pinecone(api_key=settings.PINECONE_API_KEY)
|
|
||||||
return pc.Index(settings.PINECONE_INDEX)
|
|
||||||
|
|
||||||
# ── Qdrant helpers ────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def _qdrant_client(self) -> Any:
|
|
||||||
return QdrantClient(
|
|
||||||
url=settings.QDRANT_URL,
|
|
||||||
api_key=settings.QDRANT_API_KEY or None,
|
|
||||||
)
|
|
||||||
|
|
||||||
# ── Public API ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
async def upsert(self, user_id: str, vectors: list[VectorItem]) -> None:
|
|
||||||
"""Store encrypted vectors in the backend.
|
|
||||||
|
|
||||||
Each ``VectorItem.blob`` is base64-encoded and kept in metadata/payload
|
|
||||||
so it can be returned verbatim during search.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
user_id: Used as Pinecone namespace or Qdrant payload field.
|
|
||||||
vectors: List of encrypted vector items from the client.
|
|
||||||
"""
|
|
||||||
if self._use_pinecone():
|
|
||||||
await self._pinecone_upsert(user_id, vectors)
|
|
||||||
else:
|
|
||||||
await self._qdrant_upsert(user_id, vectors)
|
|
||||||
|
|
||||||
async def search(
|
|
||||||
self,
|
|
||||||
user_id: str,
|
|
||||||
query_blob: bytes,
|
|
||||||
top_k: int,
|
|
||||||
) -> list[VectorSearchResult]:
|
|
||||||
"""Query the vector store and return encrypted result blobs.
|
|
||||||
|
|
||||||
The query vector is derived from *query_blob* using the same
|
|
||||||
deterministic mapping as upsert.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
user_id: Scopes the search to this user's namespace.
|
|
||||||
query_blob: Encrypted query from the client.
|
|
||||||
top_k: Maximum number of results to return.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
List of ``VectorSearchResult`` with ``id``, ``score``, and ``blob``.
|
|
||||||
"""
|
|
||||||
if self._use_pinecone():
|
|
||||||
return await self._pinecone_search(user_id, query_blob, top_k)
|
|
||||||
return await self._qdrant_search(user_id, query_blob, top_k)
|
|
||||||
|
|
||||||
async def delete(self, user_id: str, vector_ids: list[str]) -> None:
|
|
||||||
"""Remove vectors by ID, scoped to *user_id*.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
user_id: Namespace / payload filter to prevent cross-user deletion.
|
|
||||||
vector_ids: List of vector IDs to remove.
|
|
||||||
"""
|
|
||||||
if self._use_pinecone():
|
|
||||||
await self._pinecone_delete(user_id, vector_ids)
|
|
||||||
else:
|
|
||||||
await self._qdrant_delete(user_id, vector_ids)
|
|
||||||
|
|
||||||
# ── Pinecone implementation ───────────────────────────────────────
|
|
||||||
|
|
||||||
async def _pinecone_upsert(self, user_id: str, vectors: list[VectorItem]) -> None:
|
|
||||||
index = self._pinecone_index()
|
|
||||||
records = [
|
|
||||||
{
|
|
||||||
"id": v.id,
|
|
||||||
"values": _blob_to_vector(v.blob),
|
|
||||||
"metadata": {
|
|
||||||
"blob": base64.b64encode(v.blob).decode(),
|
|
||||||
"checksum": v.checksum,
|
|
||||||
"user_id": user_id,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
for v in vectors
|
|
||||||
]
|
|
||||||
index.upsert(vectors=records, namespace=user_id)
|
|
||||||
|
|
||||||
async def _pinecone_search(
|
|
||||||
self, user_id: str, query_blob: bytes, top_k: int
|
|
||||||
) -> list[VectorSearchResult]:
|
|
||||||
index = self._pinecone_index()
|
|
||||||
query_vector = _blob_to_vector(query_blob)
|
|
||||||
response = index.query(
|
|
||||||
vector=query_vector,
|
|
||||||
top_k=top_k,
|
|
||||||
namespace=user_id,
|
|
||||||
include_metadata=True,
|
|
||||||
)
|
|
||||||
results: list[VectorSearchResult] = []
|
|
||||||
for match in response.get("matches", []):
|
|
||||||
blob_bytes = base64.b64decode(match["metadata"]["blob"])
|
|
||||||
results.append(
|
|
||||||
VectorSearchResult(
|
|
||||||
id=match["id"],
|
|
||||||
score=match["score"],
|
|
||||||
blob=blob_bytes,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return results
|
|
||||||
|
|
||||||
async def _pinecone_delete(self, user_id: str, vector_ids: list[str]) -> None:
|
|
||||||
index = self._pinecone_index()
|
|
||||||
index.delete(ids=vector_ids, namespace=user_id)
|
|
||||||
|
|
||||||
# ── Qdrant implementation ─────────────────────────────────────────
|
|
||||||
|
|
||||||
async def _qdrant_upsert(self, user_id: str, vectors: list[VectorItem]) -> None:
|
|
||||||
client = self._qdrant_client()
|
|
||||||
points = [
|
|
||||||
PointStruct(
|
|
||||||
id=v.id,
|
|
||||||
vector=_blob_to_vector(v.blob),
|
|
||||||
payload={
|
|
||||||
"blob": base64.b64encode(v.blob).decode(),
|
|
||||||
"checksum": v.checksum,
|
|
||||||
"user_id": user_id,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
for v in vectors
|
|
||||||
]
|
|
||||||
client.upsert(collection_name=_QDRANT_COLLECTION, points=points)
|
|
||||||
|
|
||||||
async def _qdrant_search(
|
|
||||||
self, user_id: str, query_blob: bytes, top_k: int
|
|
||||||
) -> list[VectorSearchResult]:
|
|
||||||
client = self._qdrant_client()
|
|
||||||
query_vector = _blob_to_vector(query_blob)
|
|
||||||
hits = client.search(
|
|
||||||
collection_name=_QDRANT_COLLECTION,
|
|
||||||
query_vector=query_vector,
|
|
||||||
query_filter=Filter(
|
|
||||||
must=[FieldCondition(key="user_id", match=MatchValue(value=user_id))]
|
|
||||||
),
|
|
||||||
limit=top_k,
|
|
||||||
)
|
|
||||||
return [
|
|
||||||
VectorSearchResult(
|
|
||||||
id=str(hit.id),
|
|
||||||
score=hit.score,
|
|
||||||
blob=base64.b64decode(hit.payload["blob"]),
|
|
||||||
)
|
|
||||||
for hit in hits
|
|
||||||
]
|
|
||||||
|
|
||||||
async def _qdrant_delete(self, user_id: str, vector_ids: list[str]) -> None:
|
|
||||||
client = self._qdrant_client()
|
|
||||||
client.delete(
|
|
||||||
collection_name=_QDRANT_COLLECTION,
|
|
||||||
points_selector=PointIdsList(points=vector_ids),
|
|
||||||
)
|
|
||||||
@@ -7,18 +7,21 @@ services:
|
|||||||
- path: .env
|
- path: .env
|
||||||
required: false
|
required: false
|
||||||
environment:
|
environment:
|
||||||
DATABASE_URL: postgresql+asyncpg://postgres:postgres@db:5432/adiuva
|
DATABASE_URL: postgresql+asyncpg://postgres:postgres@db:5432/adiuvai
|
||||||
|
GITHUB_COPILOT_TOKEN_DIR: /root/.config/litellm/github_copilot
|
||||||
|
volumes:
|
||||||
|
- copilot_tokens:/root/.config/litellm/github_copilot
|
||||||
depends_on:
|
depends_on:
|
||||||
db:
|
db:
|
||||||
condition: service_healthy
|
condition: service_healthy
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
|
|
||||||
db:
|
db:
|
||||||
image: postgres:16-alpine
|
image: pgvector/pgvector:pg16
|
||||||
environment:
|
environment:
|
||||||
POSTGRES_USER: postgres
|
POSTGRES_USER: postgres
|
||||||
POSTGRES_PASSWORD: postgres
|
POSTGRES_PASSWORD: postgres
|
||||||
POSTGRES_DB: adiuva
|
POSTGRES_DB: adiuvai
|
||||||
volumes:
|
volumes:
|
||||||
- postgres_data:/var/lib/postgresql/data
|
- postgres_data:/var/lib/postgresql/data
|
||||||
healthcheck:
|
healthcheck:
|
||||||
@@ -33,36 +36,6 @@ services:
|
|||||||
# image: redis:7-alpine
|
# image: redis:7-alpine
|
||||||
# restart: unless-stopped
|
# restart: unless-stopped
|
||||||
|
|
||||||
# ── Local S3-compatible storage (MinIO) ──
|
|
||||||
minio:
|
|
||||||
image: minio/minio:latest
|
|
||||||
command: server /data --console-address ":9001"
|
|
||||||
ports:
|
|
||||||
- "9000:9000"
|
|
||||||
- "9001:9001"
|
|
||||||
environment:
|
|
||||||
MINIO_ROOT_USER: minioadmin
|
|
||||||
MINIO_ROOT_PASSWORD: minioadmin
|
|
||||||
volumes:
|
|
||||||
- minio_data:/data
|
|
||||||
healthcheck:
|
|
||||||
test: ["CMD", "mc", "ready", "local"]
|
|
||||||
interval: 5s
|
|
||||||
timeout: 5s
|
|
||||||
retries: 5
|
|
||||||
restart: unless-stopped
|
|
||||||
|
|
||||||
# ── Local vector store (Qdrant) ──
|
|
||||||
qdrant:
|
|
||||||
image: qdrant/qdrant:latest
|
|
||||||
ports:
|
|
||||||
- "6333:6333"
|
|
||||||
- "6334:6334"
|
|
||||||
volumes:
|
|
||||||
- qdrant_data:/qdrant/storage
|
|
||||||
restart: unless-stopped
|
|
||||||
|
|
||||||
volumes:
|
volumes:
|
||||||
postgres_data:
|
postgres_data:
|
||||||
minio_data:
|
copilot_tokens:
|
||||||
qdrant_data:
|
|
||||||
|
|||||||
56
logging.conf
Normal file
56
logging.conf
Normal file
@@ -0,0 +1,56 @@
|
|||||||
|
[loggers]
|
||||||
|
keys=root,uvicorn,uvicorn.error,uvicorn.access,sqlalchemy,watchfiles
|
||||||
|
|
||||||
|
[handlers]
|
||||||
|
keys=console,file
|
||||||
|
|
||||||
|
[formatters]
|
||||||
|
keys=default
|
||||||
|
|
||||||
|
[logger_root]
|
||||||
|
level=INFO
|
||||||
|
handlers=console,file
|
||||||
|
|
||||||
|
[logger_uvicorn]
|
||||||
|
level=INFO
|
||||||
|
handlers=
|
||||||
|
qualname=uvicorn
|
||||||
|
propagate=1
|
||||||
|
|
||||||
|
[logger_uvicorn.error]
|
||||||
|
level=INFO
|
||||||
|
handlers=
|
||||||
|
qualname=uvicorn.error
|
||||||
|
propagate=1
|
||||||
|
|
||||||
|
[logger_uvicorn.access]
|
||||||
|
level=INFO
|
||||||
|
handlers=
|
||||||
|
qualname=uvicorn.access
|
||||||
|
propagate=1
|
||||||
|
|
||||||
|
[logger_sqlalchemy]
|
||||||
|
level=WARNING
|
||||||
|
handlers=
|
||||||
|
qualname=sqlalchemy
|
||||||
|
propagate=1
|
||||||
|
|
||||||
|
[logger_watchfiles]
|
||||||
|
level=WARNING
|
||||||
|
handlers=
|
||||||
|
qualname=watchfiles
|
||||||
|
propagate=1
|
||||||
|
|
||||||
|
[handler_console]
|
||||||
|
class=StreamHandler
|
||||||
|
formatter=default
|
||||||
|
args=(sys.stderr,)
|
||||||
|
|
||||||
|
[handler_file]
|
||||||
|
class=logging.handlers.RotatingFileHandler
|
||||||
|
formatter=default
|
||||||
|
args=('logs/app.log', 'a', 10485760, 5, 'utf-8')
|
||||||
|
|
||||||
|
[formatter_default]
|
||||||
|
format=%(asctime)s %(levelname)s %(name)s: %(message)s
|
||||||
|
datefmt=%Y-%m-%d %H:%M:%S
|
||||||
@@ -3,6 +3,7 @@ uvicorn[standard]>=0.34.0
|
|||||||
gunicorn>=22.0.0
|
gunicorn>=22.0.0
|
||||||
langchain>=0.3.0
|
langchain>=0.3.0
|
||||||
langchain-openai>=0.3.0
|
langchain-openai>=0.3.0
|
||||||
|
langchain-litellm>=0.1.0
|
||||||
litellm>=1.50.0
|
litellm>=1.50.0
|
||||||
pydantic>=2.10.0
|
pydantic>=2.10.0
|
||||||
pydantic-settings>=2.7.0
|
pydantic-settings>=2.7.0
|
||||||
@@ -24,4 +25,19 @@ aiosqlite>=0.20.0
|
|||||||
moto[s3]>=5.0.0
|
moto[s3]>=5.0.0
|
||||||
pinecone>=5.0.0
|
pinecone>=5.0.0
|
||||||
qdrant-client>=1.7.0
|
qdrant-client>=1.7.0
|
||||||
|
croniter>=3.0.0
|
||||||
|
google-api-python-client>=2.130.0
|
||||||
|
google-auth>=2.29.0
|
||||||
|
google-auth-oauthlib>=1.2.0
|
||||||
|
google-auth-httplib2>=0.2.0
|
||||||
|
msal>=1.28.0
|
||||||
|
cryptography>=42.0.0
|
||||||
|
pgvector>=0.2.5
|
||||||
|
langfuse>=3.3.1
|
||||||
|
beautifulsoup4>=4.12.0
|
||||||
|
lxml>=5.0.0
|
||||||
|
PyYAML>=6.0.0
|
||||||
|
apscheduler>=3.10.0
|
||||||
ruff>=0.8.0
|
ruff>=0.8.0
|
||||||
|
pypdf>=4.0
|
||||||
|
python-docx>=1.1
|
||||||
|
|||||||
1
results.xml
Normal file
1
results.xml
Normal file
File diff suppressed because one or more lines are too long
@@ -6,26 +6,23 @@ a per-test session, and a FastAPI ``TestClient`` wired to use it.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
import time
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import AsyncGenerator, Generator
|
from collections.abc import AsyncGenerator, Generator
|
||||||
from unittest.mock import patch
|
|
||||||
|
|
||||||
import boto3
|
|
||||||
import pytest
|
import pytest
|
||||||
import pytest_asyncio
|
import pytest_asyncio
|
||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
from jose import jwt
|
from jose import jwt
|
||||||
from moto import mock_aws
|
|
||||||
from sqlalchemy import StaticPool, event
|
from sqlalchemy import StaticPool, event
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||||
|
|
||||||
|
from sqlalchemy import select
|
||||||
|
|
||||||
from app.config.settings import settings
|
from app.config.settings import settings
|
||||||
from app.db import Base, get_session
|
from app.db import Base, get_session
|
||||||
from app.main import app
|
from app.main import app
|
||||||
from app.models import Plugin, Subscription, User
|
from app.models import Subscription, User
|
||||||
|
|
||||||
# ── Fixed test user IDs (one per tier) ───────────────────────────────
|
# ── Fixed test user IDs (one per tier) ───────────────────────────────
|
||||||
|
|
||||||
@@ -109,79 +106,6 @@ def client(db_session: AsyncSession) -> Generator[TestClient, None, None]: # n
|
|||||||
app.dependency_overrides.pop(get_session, None)
|
app.dependency_overrides.pop(get_session, None)
|
||||||
|
|
||||||
|
|
||||||
# ── Seed data helpers ────────────────────────────────────────────────
|
|
||||||
|
|
||||||
_SEED_PLUGINS = [
|
|
||||||
Plugin(
|
|
||||||
id="plugin-github-sync",
|
|
||||||
name="GitHub Sync",
|
|
||||||
description="Sync tasks with GitHub Issues and pull requests.",
|
|
||||||
version="1.0.0",
|
|
||||||
author_name="Adiuva",
|
|
||||||
category="productivity",
|
|
||||||
price_cents=0,
|
|
||||||
permissions=json.dumps(["read:tasks", "write:tasks"]),
|
|
||||||
status="approved",
|
|
||||||
s3_package_key="plugins/plugin-github-sync/1.0.0/package.zip",
|
|
||||||
install_count=0,
|
|
||||||
avg_rating=0.0,
|
|
||||||
),
|
|
||||||
Plugin(
|
|
||||||
id="plugin-slack-notify",
|
|
||||||
name="Slack Notifier",
|
|
||||||
description="Post task and checkpoint updates to Slack channels.",
|
|
||||||
version="1.2.0",
|
|
||||||
author_name="Adiuva",
|
|
||||||
category="communication",
|
|
||||||
price_cents=499,
|
|
||||||
permissions=json.dumps(["read:tasks", "read:checkpoints"]),
|
|
||||||
status="approved",
|
|
||||||
s3_package_key="plugins/plugin-slack-notify/1.2.0/package.zip",
|
|
||||||
install_count=0,
|
|
||||||
avg_rating=0.0,
|
|
||||||
),
|
|
||||||
Plugin(
|
|
||||||
id="plugin-time-tracker",
|
|
||||||
name="Time Tracker",
|
|
||||||
description="Track time spent on tasks with automatic reporting.",
|
|
||||||
version="0.9.1",
|
|
||||||
author_name="Third Party",
|
|
||||||
category="productivity",
|
|
||||||
price_cents=999,
|
|
||||||
permissions=json.dumps(["read:tasks", "write:tasks"]),
|
|
||||||
status="approved",
|
|
||||||
s3_package_key="plugins/plugin-time-tracker/0.9.1/package.zip",
|
|
||||||
install_count=0,
|
|
||||||
avg_rating=0.0,
|
|
||||||
),
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
@pytest_asyncio.fixture
|
|
||||||
async def seed_plugins(db_session: AsyncSession) -> list[Plugin]:
|
|
||||||
"""Insert the 3 default approved plugins and return them."""
|
|
||||||
plugins = []
|
|
||||||
for template in _SEED_PLUGINS:
|
|
||||||
p = Plugin(
|
|
||||||
id=template.id,
|
|
||||||
name=template.name,
|
|
||||||
description=template.description,
|
|
||||||
version=template.version,
|
|
||||||
author_name=template.author_name,
|
|
||||||
category=template.category,
|
|
||||||
price_cents=template.price_cents,
|
|
||||||
permissions=template.permissions,
|
|
||||||
status=template.status,
|
|
||||||
s3_package_key=template.s3_package_key,
|
|
||||||
install_count=template.install_count,
|
|
||||||
avg_rating=template.avg_rating,
|
|
||||||
)
|
|
||||||
db_session.add(p)
|
|
||||||
plugins.append(p)
|
|
||||||
await db_session.commit()
|
|
||||||
return plugins
|
|
||||||
|
|
||||||
|
|
||||||
# ── JWT helpers ──────────────────────────────────────────────────────
|
# ── JWT helpers ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@@ -212,24 +136,53 @@ def auth_header(tier: str = "power", user_id: str | None = None) -> dict[str, st
|
|||||||
return {"Authorization": f"Bearer {make_jwt(tier, user_id)}"}
|
return {"Authorization": f"Bearer {make_jwt(tier, user_id)}"}
|
||||||
|
|
||||||
|
|
||||||
# ── S3 mock fixture ──────────────────────────────────────────────────
|
# ── Convenience aliases and per-tier user fixtures ────────────────────
|
||||||
|
|
||||||
S3_TEST_BUCKET = "test-bucket"
|
@pytest_asyncio.fixture
|
||||||
S3_TEST_REGION = "us-east-1"
|
async def db(db_session: AsyncSession) -> AsyncSession:
|
||||||
|
"""Alias for db_session — used by folder quota tests."""
|
||||||
|
return db_session
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_user_free(db_session: AsyncSession):
|
||||||
|
"""Return the seeded free-tier User row."""
|
||||||
|
result = await db_session.execute(
|
||||||
|
select(User).where(User.id == TEST_USER_IDS["free"])
|
||||||
|
)
|
||||||
|
return result.scalar_one()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_user_power(db_session: AsyncSession):
|
||||||
|
"""Return the seeded power-tier User row."""
|
||||||
|
result = await db_session.execute(
|
||||||
|
select(User).where(User.id == TEST_USER_IDS["power"])
|
||||||
|
)
|
||||||
|
return result.scalar_one()
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def s3_bucket():
|
def auth_headers_free() -> dict[str, str]:
|
||||||
"""Create a mocked S3 bucket via moto and patch BlobStore settings."""
|
"""Authorization header for the seeded free-tier user."""
|
||||||
with mock_aws():
|
return auth_header("free")
|
||||||
os.environ.setdefault("AWS_ACCESS_KEY_ID", "testing")
|
|
||||||
os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "testing")
|
|
||||||
os.environ.setdefault("AWS_DEFAULT_REGION", S3_TEST_REGION)
|
# ── CLI options ───────────────────────────────────────────────────────
|
||||||
client = boto3.client("s3", region_name=S3_TEST_REGION)
|
|
||||||
client.create_bucket(Bucket=S3_TEST_BUCKET)
|
def pytest_addoption(parser):
|
||||||
with patch("app.storage.blob_store.settings") as mock_settings:
|
parser.addoption(
|
||||||
mock_settings.S3_BUCKET = S3_TEST_BUCKET
|
"--preprocess-dir",
|
||||||
mock_settings.S3_REGION = S3_TEST_REGION
|
default=None,
|
||||||
mock_settings.AWS_ACCESS_KEY_ID = "testing"
|
help="Override fixture folder for preprocessor tests (must contain cases.yaml + data/)",
|
||||||
mock_settings.AWS_SECRET_ACCESS_KEY = "testing"
|
)
|
||||||
yield S3_TEST_BUCKET
|
parser.addoption(
|
||||||
|
"--runner-dir",
|
||||||
|
default=None,
|
||||||
|
help="Override fixture folder for agent_runner_v2 eval tests (must contain cases.yaml + data/)",
|
||||||
|
)
|
||||||
|
parser.addoption(
|
||||||
|
"--journey-dir",
|
||||||
|
default=None,
|
||||||
|
help="Override fixture folder for journey_v2 eval tests (must contain cases.yaml + data/)",
|
||||||
|
)
|
||||||
|
|||||||
86
tests/fixtures/agent_runner_v2/cases.yaml
vendored
Normal file
86
tests/fixtures/agent_runner_v2/cases.yaml
vendored
Normal file
@@ -0,0 +1,86 @@
|
|||||||
|
# Agent Runner V2 — eval test cases (Step 2, requires real LLM)
|
||||||
|
#
|
||||||
|
# Each case drives one parametrized `test_eval_runner` invocation.
|
||||||
|
#
|
||||||
|
# Keys
|
||||||
|
# ----
|
||||||
|
# id: str unique identifier shown in pytest output
|
||||||
|
# description: str human-readable label
|
||||||
|
# file: str filename inside data/
|
||||||
|
# file_path: str path reported to the executor (affects project-matching via filename)
|
||||||
|
# projects: [alpha|beta] symbolic project names resolved by the test helper
|
||||||
|
#
|
||||||
|
# Optional pre-existing records (dedup tests)
|
||||||
|
# existing_tasks: list of {id, title, status, priority}
|
||||||
|
# existing_notes: list of {id, title, content}
|
||||||
|
# existing_timelines: list of {id, title, date}
|
||||||
|
#
|
||||||
|
# Assertions (one or more)
|
||||||
|
# expect_insert: <table> at least 1 insert row in this table (tasks|notes|timelines)
|
||||||
|
# expect_no_insert: true zero inserts in any table
|
||||||
|
# expect_project_id: <id> any insert must carry this projectId
|
||||||
|
# expect_dedup: true task inserts == 0 OR task updates >= 1 (dedup check)
|
||||||
|
#
|
||||||
|
# Langfuse
|
||||||
|
# score_name: str observation score name
|
||||||
|
|
||||||
|
- id: "2.1"
|
||||||
|
description: "Action email → create_task"
|
||||||
|
file: email_action.html
|
||||||
|
file_path: /emails/ProjectAlpha_action.html
|
||||||
|
projects: [alpha, beta]
|
||||||
|
expect_insert: tasks
|
||||||
|
score_name: runner.email_to_task
|
||||||
|
|
||||||
|
- id: "2.2"
|
||||||
|
description: "Informational email → create_note"
|
||||||
|
file: email_info.html
|
||||||
|
file_path: /emails/ProjectAlpha_info.html
|
||||||
|
projects: [alpha, beta]
|
||||||
|
expect_insert: notes
|
||||||
|
score_name: runner.email_to_note
|
||||||
|
|
||||||
|
- id: "2.3"
|
||||||
|
description: "Email with meeting date → create_timeline"
|
||||||
|
file: email_date.html
|
||||||
|
file_path: /emails/ProjectAlpha_kickoff.html
|
||||||
|
projects: [alpha, beta]
|
||||||
|
expect_insert: timelines
|
||||||
|
score_name: runner.email_to_timeline
|
||||||
|
|
||||||
|
- id: "2.4"
|
||||||
|
description: "Filename contains project name → correct project assigned"
|
||||||
|
file: email_action.html
|
||||||
|
file_path: /emails/ProjectAlpha_report.html
|
||||||
|
projects: [alpha, beta]
|
||||||
|
expect_project_id: proj-alpha
|
||||||
|
score_name: runner.project_filename
|
||||||
|
|
||||||
|
- id: "2.5"
|
||||||
|
description: "Email body mentions project → correct project assigned"
|
||||||
|
file: email_action.html
|
||||||
|
file_path: /emails/email_001.html
|
||||||
|
projects: [alpha, beta]
|
||||||
|
expect_project_id: proj-alpha
|
||||||
|
score_name: runner.project_content
|
||||||
|
|
||||||
|
- id: "2.6"
|
||||||
|
description: "Newsletter + global rule no-project → no creates"
|
||||||
|
file: email_no_project.html
|
||||||
|
file_path: /emails/newsletter.html
|
||||||
|
projects: [alpha, beta]
|
||||||
|
expect_no_insert: true
|
||||||
|
score_name: runner.no_project
|
||||||
|
|
||||||
|
- id: "2.7"
|
||||||
|
description: "Existing task with same title → dedup (update not create)"
|
||||||
|
file: email_action.html
|
||||||
|
file_path: /emails/ProjectAlpha_followup.html
|
||||||
|
projects: [alpha]
|
||||||
|
existing_tasks:
|
||||||
|
- id: task-existing
|
||||||
|
title: Fix the login bug
|
||||||
|
status: todo
|
||||||
|
priority: medium
|
||||||
|
expect_dedup: true
|
||||||
|
score_name: runner.dedup
|
||||||
7
tests/fixtures/agent_runner_v2/data/email_action.html
vendored
Normal file
7
tests/fixtures/agent_runner_v2/data/email_action.html
vendored
Normal file
@@ -0,0 +1,7 @@
|
|||||||
|
<html><head></head><body>
|
||||||
|
<p><b>From:</b> boss@company.com</p>
|
||||||
|
<p><b>To:</b> dev@company.com</p>
|
||||||
|
<p><b>Subject:</b> Fix the login bug</p>
|
||||||
|
<p><b>Date:</b> 2026-04-07</p>
|
||||||
|
<p>Hi,<br>Please fix the login bug in Project Alpha by Friday. High priority!</p>
|
||||||
|
</body></html>
|
||||||
5
tests/fixtures/agent_runner_v2/data/email_date.html
vendored
Normal file
5
tests/fixtures/agent_runner_v2/data/email_date.html
vendored
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
<html><head></head><body>
|
||||||
|
<p><b>From:</b> pm@company.com</p>
|
||||||
|
<p><b>Subject:</b> Project Alpha kick-off meeting</p>
|
||||||
|
<p>The kick-off meeting for Project Alpha is scheduled for 2026-04-15 at 10:00.</p>
|
||||||
|
</body></html>
|
||||||
7
tests/fixtures/agent_runner_v2/data/email_info.html
vendored
Normal file
7
tests/fixtures/agent_runner_v2/data/email_info.html
vendored
Normal file
@@ -0,0 +1,7 @@
|
|||||||
|
<html><head></head><body>
|
||||||
|
<p><b>From:</b> pm@company.com</p>
|
||||||
|
<p><b>To:</b> team@company.com</p>
|
||||||
|
<p><b>Subject:</b> FYI: New policy for Project Alpha</p>
|
||||||
|
<p>Just a heads-up that starting next week all code reviews must be done
|
||||||
|
within 24 hours for Project Alpha. No action needed from you now.</p>
|
||||||
|
</body></html>
|
||||||
5
tests/fixtures/agent_runner_v2/data/email_no_project.html
vendored
Normal file
5
tests/fixtures/agent_runner_v2/data/email_no_project.html
vendored
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
<html><head></head><body>
|
||||||
|
<p><b>From:</b> newsletter@ads.com</p>
|
||||||
|
<p><b>Subject:</b> Weekly newsletter</p>
|
||||||
|
<p>Check out our latest deals on electronics!</p>
|
||||||
|
</body></html>
|
||||||
19
tests/fixtures/journey_v2/cases.yaml
vendored
Normal file
19
tests/fixtures/journey_v2/cases.yaml
vendored
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
# Journey V2 eval test cases — Step 4
|
||||||
|
#
|
||||||
|
# Only case 4.1 is kept as an automated eval. Cases 4.2–4.5 (multi-turn
|
||||||
|
# conversations that expect the LLM to produce a complete AgentConfig)
|
||||||
|
# are non-deterministic and tested manually — results tracked in Langfuse.
|
||||||
|
#
|
||||||
|
# Assertion keys:
|
||||||
|
# expect_question: true → first reply must contain "?"
|
||||||
|
|
||||||
|
- id: "4.1"
|
||||||
|
description: "Journey start explores directory, first reply contains a question"
|
||||||
|
directory: "/test/emails"
|
||||||
|
data_types: ["tasks", "notes", "timelines"]
|
||||||
|
directory_files:
|
||||||
|
- path: "/test/emails/outlook_export_2024.html"
|
||||||
|
content_file: "email_action.html"
|
||||||
|
user_messages: []
|
||||||
|
score_name: "journey.start"
|
||||||
|
expect_question: true
|
||||||
23
tests/fixtures/journey_v2/data/email_action.html
vendored
Normal file
23
tests/fixtures/journey_v2/data/email_action.html
vendored
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html>
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<title>Email: Fix the login bug</title>
|
||||||
|
<style>body { font-family: Arial; } .header { color: #666; }</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="header">
|
||||||
|
<p><strong>From:</strong> boss@company.com</p>
|
||||||
|
<p><strong>To:</strong> dev@company.com</p>
|
||||||
|
<p><strong>Subject:</strong> Fix the login bug</p>
|
||||||
|
<p><strong>Date:</strong> Mon, 7 Apr 2026 09:15:00 +0000</p>
|
||||||
|
</div>
|
||||||
|
<div class="body">
|
||||||
|
<p>Hi,</p>
|
||||||
|
<p>Please fix the login bug in Project Alpha as soon as possible.
|
||||||
|
Users are reporting that they can't log in with their Google accounts.
|
||||||
|
This is blocking the whole team. Please resolve it by Friday.</p>
|
||||||
|
<p>Thanks,<br>Boss</p>
|
||||||
|
</div>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user