Compare commits
10 Commits
f7404b6f66
...
feature/de
| Author | SHA1 | Date | |
|---|---|---|---|
| 47bf1881e5 | |||
| 24a9c1b752 | |||
| 706bf88883 | |||
| 4ff0b27084 | |||
| 61d2a18234 | |||
| b3687719b6 | |||
| f80bdfa8f7 | |||
| 617a17db40 | |||
| 92716cb89a | |||
| cfc9d7a942 |
@@ -1,523 +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) |
|
||||
| `timelines` | 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/timeline_agent.py` (4 tools):**
|
||||
- `list_timelines(project_id)`: `execute_on_client(action="select", table="timelines", filters={projectId})` → format + return
|
||||
- `create_timeline(project_id, title, date, ...)`: `execute_on_client(action="insert", table="timelines", data={...})` → return confirmation + id
|
||||
- `update_timeline(timeline_id, ...)`: build updates → `execute_on_client(action="update", ...)` → return confirmation
|
||||
- `delete_timeline(timeline_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/timeline_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.
|
||||
|
||||
---
|
||||
|
||||
## Phase 3 — Agent System: Config, Orchestration & Cloud Connectors
|
||||
|
||||
> **Objective:** Backend manages all agent configuration, scheduling, orchestration, and cloud data fetching. Two agent types: **Local Directory Agent** (backend triggers Electron to read files, then AI analyzes) and **Cloud Connector Agent** (backend fetches Gmail/Teams data directly, AI analyzes, pushes results to Electron via WS tool_call). All extracted items use existing WS tool infrastructure to insert into Electron's local DB with `is_ai_suggested=True`.
|
||||
>
|
||||
> **Electron Phase 3 plan:** `../adiuva/AI_REFACTOR_PLAN.md` Phase 3 section.
|
||||
>
|
||||
> **Electron UI status (2025):** Steps 3.6, 3.7, 3.8 of the Electron plan are ✅ complete. Agents are configured inside the Settings page (`/settings?section=agents`) — not a standalone route. The `JourneyDialog` (Step 3.8) is embedded inline in the Settings → Agents section. `LocalAgentConfigPanel` and `CloudAgentConfigPanel` (Step 3.7) are also inline. This affects the journey API contract (see Step 3.5 below).
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
Local Agent:
|
||||
Scheduler/manual trigger ──► check device online ──► WS agent_run → Electron
|
||||
──► Electron reads files ──► WS agent_data → Backend
|
||||
──► Backend AI (prompt_template + file content) ──► WS tool_call(insert) → Electron
|
||||
──► Electron persists with isAiSuggested=1
|
||||
|
||||
Cloud Agent:
|
||||
Scheduler/manual trigger ──► Backend fetches Gmail/Teams (OAuth) ──► Backend AI analyzes
|
||||
──► check device online ──► WS tool_call(insert) → Electron ──► Electron persists
|
||||
```
|
||||
|
||||
**New WS frame types:**
|
||||
|
||||
| Direction | `type` | Payload |
|
||||
|---|---|---|
|
||||
| Server → Client | `agent_run` | `{ run_id, agent_id, config: { paths, file_extensions, prompt_template, data_types } }` |
|
||||
| Client → Server | `agent_data` | `{ run_id, files: [{ path, name, content, metadata }] }` |
|
||||
| Client → Server | `agent_complete` | `{ run_id, files_read, errors }` |
|
||||
| Client → Server | `device_hello` | `{ device_id, agent_ids }` |
|
||||
|
||||
### Step 3.1 — Agent config tables
|
||||
- [x] Add to `app/models.py`:
|
||||
- **`LocalAgentConfig`**:
|
||||
- `id` UUID PK
|
||||
- `user_id` FK → users
|
||||
- `device_id` str — identifies which Electron install this config belongs to
|
||||
- `name` str
|
||||
- `directory_paths` JSON — list of absolute paths on the device
|
||||
- `data_types` JSON — which tables to extract to: `["tasks", "notes", "timelines", "projects"]`
|
||||
- `prompt_template` text — user-configured via Chatbot Journey
|
||||
- `file_extensions` JSON — e.g. `[".eml", ".txt", ".pdf", ".md"]`
|
||||
- `schedule_cron` str — e.g. `"0 */6 * * *"` (every 6h)
|
||||
- `enabled` bool (default True)
|
||||
- `last_run_at` datetime nullable
|
||||
- `created_at`, `updated_at` timestamps
|
||||
- **`CloudAgentConfig`**:
|
||||
- `id` UUID PK
|
||||
- `user_id` FK → users
|
||||
- `provider` str — enum: `gmail`, `teams`, `outlook`
|
||||
- `name` str
|
||||
- `data_types` JSON — same format as local
|
||||
- `prompt_template` text
|
||||
- `oauth_token_encrypted` text — Fernet-encrypted OAuth2 credentials
|
||||
- `schedule_cron` str
|
||||
- `enabled` bool (default True)
|
||||
- `last_run_at` datetime nullable
|
||||
- `filter_config` JSON — provider-specific: `{ labels: [], date_range: {from, to}, senders: [] }`
|
||||
- `created_at`, `updated_at` timestamps
|
||||
- **`AgentRunLog`**:
|
||||
- `id` UUID PK
|
||||
- `agent_id` str — references LocalAgentConfig.id or CloudAgentConfig.id
|
||||
- `agent_type` str — `local` or `cloud`
|
||||
- `user_id` FK → users
|
||||
- `status` str — `running`, `success`, `error`, `partial`
|
||||
- `items_processed` int (default 0)
|
||||
- `items_created` int (default 0)
|
||||
- `errors` JSON — list of error strings
|
||||
- `started_at` datetime
|
||||
- `completed_at` datetime nullable
|
||||
- [x] Add Pydantic schemas to `app/schemas.py`:
|
||||
- `LocalAgentConfigCreate`, `LocalAgentConfigUpdate`, `LocalAgentConfigResponse`
|
||||
- `CloudAgentConfigCreate`, `CloudAgentConfigUpdate`, `CloudAgentConfigResponse`
|
||||
- `AgentRunLogResponse`
|
||||
- `AgentCatalogItem` — `{ type, name, description, config_schema }`
|
||||
- `WsAgentRun`, `WsAgentData`, `WsAgentComplete`, `WsDeviceHello`
|
||||
- [x] Generate Alembic migration
|
||||
- **Files:** `app/models.py`, `app/schemas.py`, `alembic/versions/`
|
||||
- **Outcome:** Agent config and run tracking tables in PostgreSQL.
|
||||
|
||||
### Step 3.2 — Agent CRUD API routes
|
||||
- [x] Create `app/api/routes/agents.py`:
|
||||
- `GET /api/v1/agents/catalog` — returns hardcoded agent type catalog:
|
||||
- `local_directory`: "Watches local directories, extracts data from files using AI"
|
||||
- `gmail`: "Scans Gmail inbox, extracts tasks/notes from emails"
|
||||
- `teams`: "Monitors Teams messages, extracts action items"
|
||||
- `outlook`: "Scans Outlook inbox, extracts tasks/notes"
|
||||
- `GET /api/v1/agents/local` — list user's local agent configs
|
||||
- `POST /api/v1/agents/local` — create local agent config
|
||||
- Body: `{ name, device_id, directory_paths, data_types, prompt_template, file_extensions, schedule_cron }`
|
||||
- Tier check: count enabled agents ≤ `batch_active` limit
|
||||
- `PUT /api/v1/agents/local/{id}` — update config (ownership check)
|
||||
- `DELETE /api/v1/agents/local/{id}` — delete config + associated run logs
|
||||
- `GET /api/v1/agents/cloud` — list user's cloud agent configs
|
||||
- `POST /api/v1/agents/cloud` — create cloud connector config
|
||||
- Body: `{ provider, name, data_types, prompt_template, oauth_token_encrypted, schedule_cron, filter_config }`
|
||||
- Tier check: same `batch_active` limit (local + cloud count together)
|
||||
- `PUT /api/v1/agents/cloud/{id}` — update config
|
||||
- `DELETE /api/v1/agents/cloud/{id}` — delete config + run logs
|
||||
- `GET /api/v1/agents/runs` — query params: `agent_id`, `page`, `limit` → paginated run logs
|
||||
- `POST /api/v1/agents/{id}/run` — manual trigger (dispatches to agent runner)
|
||||
- All routes require JWT auth; ownership enforced on all mutations
|
||||
- [x] Register router in `app/main.py`
|
||||
- **Files:** `app/api/routes/agents.py`, `app/main.py`
|
||||
- **Outcome:** Full CRUD for agent configs with tier-gated creation limits.
|
||||
|
||||
### Step 3.3 — Device WS endpoint
|
||||
- [x] Create `app/api/routes/device_ws.py`:
|
||||
- `WebSocket /api/v1/ws/device?token=<jwt>` — persistent connection from Electron
|
||||
- On connect:
|
||||
- Authenticate JWT
|
||||
- Receive `device_hello` frame → extract `device_id`, `agent_ids`
|
||||
- Store connection in `DeviceConnectionManager` (in-memory dict: `user_id → { ws, device_id }`)
|
||||
- Check for overdue agent runs → trigger them immediately
|
||||
- Message loop:
|
||||
- `agent_data` → route to active agent run handler
|
||||
- `agent_complete` → finalize agent run
|
||||
- `tool_result` → route to pending tool call (same pattern as chat WS)
|
||||
- `pong` → heartbeat ack
|
||||
- On disconnect:
|
||||
- Remove from `DeviceConnectionManager`
|
||||
- Mark any in-progress agent runs as `error` with "device disconnected"
|
||||
- Heartbeat: send `ping` every 30s, disconnect if no `pong` within 10s
|
||||
- [x] Create `app/core/device_manager.py`:
|
||||
- `DeviceConnectionManager` (singleton):
|
||||
- `register(user_id, device_id, ws)` — stores active connection
|
||||
- `unregister(user_id)` — removes connection
|
||||
- `get_ws(user_id) -> WebSocket | None` — returns active WS if device is online
|
||||
- `is_online(user_id, device_id=None) -> bool` — optionally checks specific device
|
||||
- `send_frame(user_id, frame: dict)` — sends JSON frame to device
|
||||
- **Files:** `app/api/routes/device_ws.py`, `app/core/device_manager.py`, `app/main.py`
|
||||
- **Outcome:** Backend maintains persistent WS connections to Electron devices for agent triggers.
|
||||
|
||||
### Step 3.4 — Agent run orchestrator
|
||||
- [x] Create `app/core/agent_runner.py`:
|
||||
- `async run_local_agent(user_id, config: LocalAgentConfig, device_mgr: DeviceConnectionManager)`:
|
||||
1. Check device is online with matching `device_id` → abort if offline
|
||||
2. Create `AgentRunLog` with `status=running`
|
||||
3. Send `WsAgentRun` frame to Electron with config (paths, extensions, prompt)
|
||||
4. Await `WsAgentData` frames — collect file contents
|
||||
5. Await `WsAgentComplete` frame — Electron signals done reading
|
||||
6. For each file: call LLM with `prompt_template` + file content → extract structured items
|
||||
7. For each extracted item: send `WsToolCall(insert, table, data)` to Electron → await `WsToolResult`
|
||||
- All inserts include `is_ai_suggested=True, is_approved=False`
|
||||
8. Update `AgentRunLog`: `status=success`, `items_processed`, `items_created`
|
||||
- `async run_cloud_agent(user_id, config: CloudAgentConfig, device_mgr: DeviceConnectionManager)`:
|
||||
1. Check device is online → abort if offline (results must push to Electron)
|
||||
2. Create `AgentRunLog` with `status=running`
|
||||
3. Decrypt OAuth credentials from `config.oauth_token_encrypted`
|
||||
4. Fetch data from cloud provider (Step 3.6):
|
||||
- Gmail: `google-api-python-client` + `filter_config` label/date filters
|
||||
- Teams: `msgraph-sdk` + channel/date filters
|
||||
- Outlook: `msgraph-sdk` + folder/date filters
|
||||
5. For each item: call LLM with `prompt_template` + email/message content → extract structured items
|
||||
6. For each extracted item: send `WsToolCall(insert)` to Electron → await `WsToolResult`
|
||||
7. Update `AgentRunLog`
|
||||
- `async trigger_pending_runs(user_id, device_id, device_mgr)`:
|
||||
- Called when Electron connects (after `device_hello`)
|
||||
- Queries all enabled agent configs where `last_run_at + schedule_interval < now()`
|
||||
- For local agents: only triggers if `config.device_id == device_id`
|
||||
- For cloud agents: triggers regardless of device (any connected device can receive results)
|
||||
- Executes runs sequentially (one at a time to avoid overwhelming the WS)
|
||||
- Error handling: on any failure, update `AgentRunLog` with `status=error` + error details
|
||||
- [x] Wire `POST /agents/{id}/run` endpoint to dispatch background task via `asyncio.create_task()`
|
||||
- [x] Replace `_trigger_pending_runs_stub` in `device_ws.py` with real `trigger_pending_runs` call
|
||||
- [x] Add `croniter>=3.0.0` to `requirements.txt`
|
||||
- [x] 23 unit + integration tests covering all code paths
|
||||
- **Files:** `app/core/agent_runner.py`, `app/api/routes/agents.py`, `app/api/routes/device_ws.py`, `requirements.txt`, `tests/test_agent_runner.py`
|
||||
- **Outcome:** Backend drives all agent execution — both local (via WS file request) and cloud (direct API calls — stub until Step 3.6).
|
||||
|
||||
### Step 3.5 — Chatbot Journey endpoint
|
||||
- [x] Create `app/api/routes/agent_setup.py`:
|
||||
- `POST /api/v1/agents/journey/start`:
|
||||
- Body: `{ agent_type: "local"|"cloud", agent_id: str | None }`
|
||||
- `agent_type`: which kind of agent this journey configures.
|
||||
- `agent_id`: optional — if provided, the session is pre-seeded with the existing agent's `prompt_template` so the user can refine it. If absent, fresh journey.
|
||||
- **No `data_types` field** — data types are determined through the conversation itself, not sent upfront.
|
||||
- Creates a journey session (in-memory or Redis-backed)
|
||||
- Returns first AI message: contextual question based on agent type
|
||||
- Local: "What kind of files are in the directories you want to monitor? (emails, documents, logs, etc.)"
|
||||
- Cloud: "What kind of emails/messages should I look for? (client communications, invoices, meeting notes, etc.)"
|
||||
- Response: `{ session_id, message, done: false }`
|
||||
- **Electron note:** `proxyPost` auto-converts camelCase keys to snake_case. Electron sends `{ agentType, agentId }` → backend receives `{ agent_type, agent_id }`.
|
||||
- `POST /api/v1/agents/journey/message`:
|
||||
- Body: `{ session_id, message }`
|
||||
- AI processes user's answer, asks follow-up questions (max 5 turns)
|
||||
- System prompt: "You are configuring a data extraction agent for a freelancer. Ask about file format, what data to extract (tasks, notes, timelines), naming conventions, priority rules, and any special mapping. After 3-5 questions, generate a detailed prompt_template."
|
||||
- When AI determines enough context: `{ session_id, message: "Here's your configuration...", done: true, prompt_template: "..." }`
|
||||
- The `prompt_template` is a structured instruction for the extraction LLM (e.g. "Extract tasks from email. Subject becomes task title. If body contains 'urgent' or 'ASAP', set priority to 'high'. Extract due dates if mentioned.")
|
||||
- **Electron note:** `toCamelCase` converts the response → Electron reads `promptTemplate` from the final message and auto-fills the agent config panel. User clicks "Save & apply" which calls `agent.local.update` / `agent.cloud.update` tRPC mutation.
|
||||
- **Files:** `app/api/routes/agent_setup.py`, `app/main.py`
|
||||
- **Outcome:** Users configure AI prompts through guided conversation. Journey can refine an existing config when `agent_id` is provided. ✅
|
||||
|
||||
### Step 3.6 — Cloud provider integrations
|
||||
- [x] Create `app/integrations/gmail.py`:
|
||||
- `GmailClient`:
|
||||
- `__init__(oauth_token)` — initializes Google API client
|
||||
- `async fetch_messages(filter_config, since: datetime) -> list[EmailMessage]`
|
||||
- `EmailMessage`: `{ id, subject, sender, body_text, date, labels }`
|
||||
- Handles token refresh via Google OAuth2 refresh flow
|
||||
- Respects `filter_config.labels`, `filter_config.date_range`, `filter_config.senders`
|
||||
- [x] Create `app/integrations/ms_graph.py`:
|
||||
- `MSGraphClient`:
|
||||
- `__init__(oauth_token)` — initializes MS Graph client
|
||||
- `async fetch_emails(filter_config, since: datetime) -> list[EmailMessage]` (Outlook)
|
||||
- `async fetch_messages(filter_config, since: datetime) -> list[ChatMessage]` (Teams)
|
||||
- `ChatMessage`: `{ id, content, sender, channel, date }`
|
||||
- Handles token refresh via MSAL
|
||||
- [x] Create `app/integrations/__init__.py` — factory: `get_provider(provider_name) -> GmailClient | MSGraphClient`
|
||||
- **Dependencies:** `google-api-python-client`, `google-auth-oauthlib`, `msgraph-sdk`, `msal`
|
||||
- **Files:** `app/integrations/gmail.py`, `app/integrations/ms_graph.py`, `app/integrations/__init__.py`
|
||||
- **Outcome:** Backend can fetch emails/messages from Gmail, Outlook, and Teams.
|
||||
|
||||
### Step 3.7 — Agent scheduler
|
||||
- [ ] Create `app/core/agent_scheduler.py`:
|
||||
- Uses `APScheduler` (or simple asyncio loop) to check agent schedules
|
||||
- Every 60s: query enabled agents where `last_run_at + cron_interval < now()`
|
||||
- For each due agent:
|
||||
- Check if user's device is online via `DeviceConnectionManager`
|
||||
- If online: dispatch to `agent_runner`
|
||||
- If offline: skip (will trigger on next `device_hello`)
|
||||
- Locks: use PostgreSQL advisory locks to prevent duplicate runs in multi-instance deployments
|
||||
- [ ] Integrate with FastAPI lifespan (start scheduler on app startup, shutdown gracefully)
|
||||
- **Dependencies:** `apscheduler>=4.0`
|
||||
- **Files:** `app/core/agent_scheduler.py`, `app/main.py`
|
||||
- **Outcome:** Agents run automatically on their configured schedules.
|
||||
|
||||
### Step 3.8 — OAuth flow endpoints
|
||||
- [ ] Create `app/api/routes/oauth.py`:
|
||||
- `GET /api/v1/oauth/{provider}/authorize` — returns OAuth authorization URL
|
||||
- Gmail: Google OAuth2 with `gmail.readonly` scope
|
||||
- Outlook/Teams: MS identity platform with `Mail.Read`, `ChannelMessage.Read.All` scopes
|
||||
- `GET /api/v1/oauth/{provider}/callback` — handles OAuth redirect
|
||||
- Exchanges auth code for access + refresh tokens
|
||||
- Encrypts tokens with Fernet (server-side key from settings)
|
||||
- Returns encrypted token blob for storage in `CloudAgentConfig.oauth_token_encrypted`
|
||||
- `POST /api/v1/oauth/{provider}/refresh` — refresh expired OAuth token
|
||||
- **Files:** `app/api/routes/oauth.py`, `app/main.py`
|
||||
- **Outcome:** Users can connect Gmail/Teams/Outlook accounts securely.
|
||||
|
||||
---
|
||||
|
||||
### Phase 3 — Verification
|
||||
|
||||
| # | Scenario | Expected |
|
||||
|---|---|---|
|
||||
| 1 | **Agent CRUD** | Create/read/update/delete local and cloud configs; tier limits enforced (free=2, pro=10) |
|
||||
| 2 | **WS device connect** | Electron connects → `device_hello` → backend stores connection → triggers overdue runs |
|
||||
| 3 | **Local agent run** | Backend sends `agent_run` → Electron reads files → `agent_data` → backend AI extracts → `tool_call(insert)` → Electron persists with `isAiSuggested=1` |
|
||||
| 4 | **Cloud agent run** | Backend fetches Gmail → AI extracts tasks → `tool_call(insert)` → Electron persists |
|
||||
| 5 | **Device binding** | Local agent config with `device_id=A` only triggers when device A is connected |
|
||||
| 6 | **Chatbot Journey** | Start journey → 3-5 Q&A turns → produces valid `prompt_template` |
|
||||
| 7 | **Schedule** | Agent with `schedule_cron="0 */6 * * *"` runs every 6h when device is online |
|
||||
| 8 | **Offline resilience** | Device offline → runs skipped → device reconnects → overdue runs trigger immediately |
|
||||
| 9 | **OAuth flow** | Gmail authorize → callback → token encrypted → stored in config → fetch emails works |
|
||||
|
||||
### Phase 3 — New Dependencies
|
||||
|
||||
| Package | Purpose |
|
||||
|---|---|
|
||||
| `google-api-python-client` | Gmail API access |
|
||||
| `google-auth-oauthlib` | Gmail OAuth2 flow |
|
||||
| `msgraph-sdk` | Outlook + Teams API access |
|
||||
| `msal` | MS identity platform auth |
|
||||
| `apscheduler>=4.0` | Agent scheduling |
|
||||
| `cryptography` (Fernet) | OAuth token encryption at rest |
|
||||
|
||||
---
|
||||
|
||||
## ~~Phase 5 — Shared Memory~~ (SUPERSEDED)
|
||||
|
||||
> **This phase has been fully replaced by `V3_MIGRATION_PLAN.md`.**
|
||||
>
|
||||
> - Chat WS fix → V3 Step 5 (Unified WS Handler — single multiplexed socket)
|
||||
> - Agent memory → V3 Steps 6–7 (Cloud-side MemGPT-style memory in PostgreSQL + pgvector, encrypted at rest with per-user Fernet key)
|
||||
>
|
||||
> The on-device KV approach (Electron SQLite `agent_memory` table) is no longer the target architecture.
|
||||
> See `V3_MIGRATION_PLAN.md` for the current plan.
|
||||
572
BACKEND_PLAN.md
572
BACKEND_PLAN.md
@@ -1,572 +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/timeline_agent.py` — `@registry.register`:
|
||||
- Description: "Manages project timelines (milestones): list, create, update, delete"
|
||||
- Tools (4): `list_timelines(project_id)`, `create_timeline(project_id, title, date, is_ai_suggested, is_approved)`, `update_timeline(timeline_id, ...)`, `delete_timeline(timeline_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, Timelines, 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}` |
|
||||
| GET | `/api/v1/agents/catalog` | JWT | — | `AgentCatalogItem[]` |
|
||||
| GET | `/api/v1/agents/local` | JWT | — | `LocalAgentConfigResponse[]` |
|
||||
| POST | `/api/v1/agents/local` | JWT | `LocalAgentConfigCreate` | `LocalAgentConfigResponse` |
|
||||
| PUT | `/api/v1/agents/local/{id}` | JWT | `LocalAgentConfigUpdate` | `LocalAgentConfigResponse` |
|
||||
| DELETE | `/api/v1/agents/local/{id}` | JWT | — | `{ok: true}` |
|
||||
| GET | `/api/v1/agents/cloud` | JWT | — | `CloudAgentConfigResponse[]` |
|
||||
| POST | `/api/v1/agents/cloud` | JWT | `CloudAgentConfigCreate` | `CloudAgentConfigResponse` |
|
||||
| PUT | `/api/v1/agents/cloud/{id}` | JWT | `CloudAgentConfigUpdate` | `CloudAgentConfigResponse` |
|
||||
| DELETE | `/api/v1/agents/cloud/{id}` | JWT | — | `{ok: true}` |
|
||||
| GET | `/api/v1/agents/runs` | JWT | `?agent_id&page&limit` | `AgentRunLogResponse[]` |
|
||||
| POST | `/api/v1/agents/{id}/run` | JWT | — | `{ok: true, run_id}` |
|
||||
| POST | `/api/v1/agents/journey/start` | JWT | `{agent_type, data_types}` | `{session_id, message, done}` |
|
||||
| POST | `/api/v1/agents/journey/message` | JWT | `{session_id, message}` | `{session_id, message, done, prompt_template?}` |
|
||||
| GET | `/api/v1/oauth/{provider}/authorize` | JWT | — | `{authorization_url}` |
|
||||
| GET | `/api/v1/oauth/{provider}/callback` | — | OAuth code | `{encrypted_token}` |
|
||||
| WS | `/api/v1/ws/device` | JWT | `device_hello` (first frame) | Agent trigger + tool_call frames |
|
||||
|
||||
---
|
||||
|
||||
## 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 |
|
||||
| Cloud integrations | google-api-python-client, msgraph-sdk, msal |
|
||||
| Agent scheduling | APScheduler |
|
||||
| Testing | pytest + pytest-asyncio + httpx + moto (S3 mock) |
|
||||
| Deployment | Docker → fly.io / Railway / AWS ECS |
|
||||
|
||||
---
|
||||
|
||||
## Phase 3 — New Files
|
||||
|
||||
| File | Purpose |
|
||||
|---|---|
|
||||
| `app/models.py` | Add `LocalAgentConfig`, `CloudAgentConfig`, `AgentRunLog` models |
|
||||
| `app/schemas.py` | Add agent config schemas + WS agent frame types |
|
||||
| `app/api/routes/agents.py` | Agent CRUD endpoints (catalog, local, cloud, runs, manual trigger) |
|
||||
| `app/api/routes/agent_setup.py` | Chatbot Journey endpoints (start + message) |
|
||||
| `app/api/routes/device_ws.py` | Persistent device WS endpoint (`/api/v1/ws/device`) |
|
||||
| `app/api/routes/oauth.py` | OAuth authorize/callback for Gmail, Teams, Outlook |
|
||||
| `app/core/agent_runner.py` | Agent run orchestration — local (WS file request) + cloud (API fetch) |
|
||||
| `app/core/device_manager.py` | `DeviceConnectionManager` — tracks active Electron WS connections |
|
||||
| `app/core/agent_scheduler.py` | Periodic scheduler for agent cron triggers |
|
||||
| `app/integrations/gmail.py` | Gmail API client (fetch messages with filters) |
|
||||
| `app/integrations/ms_graph.py` | MS Graph client for Outlook emails + Teams messages |
|
||||
| `app/integrations/__init__.py` | Provider factory |
|
||||
|
||||
> **Full Phase 3 step-by-step plan:** See `AI_REFACTOR_PLAN.md` Phase 3 section.
|
||||
|
||||
---
|
||||
|
||||
## 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>`.
|
||||
@@ -1,353 +0,0 @@
|
||||
# V3 Migration Plan — Multi-Agent AI Productivity App
|
||||
|
||||
> Incremental migration from current architecture to v3.
|
||||
> Each step is self-contained, testable, and backwards-compatible.
|
||||
> No BYOK — server manages all LLM keys.
|
||||
> Memory encryption: server-side per-user Fernet key (Option A).
|
||||
|
||||
---
|
||||
|
||||
## General Rules
|
||||
|
||||
**Code Cleanup**: As you implement each step, remove any code that becomes unused or obsolete. This includes:
|
||||
- Old functions/methods that are superseded by new ones
|
||||
- Deprecated imports or modules
|
||||
- Dead code paths
|
||||
- Old test files no longer needed
|
||||
|
||||
This keeps the codebase clean and prevents confusion. When removing code, note it in the commit message if significant.
|
||||
|
||||
---
|
||||
|
||||
## Decisions Log
|
||||
|
||||
| Topic | Decision |
|
||||
|---|---|
|
||||
| WS topology | Single multiplexed socket (merge chat into device WS) |
|
||||
| LLM keys | Server-managed only, no user key passthrough |
|
||||
| Memory encryption | Per-user server-generated Fernet key, encrypted at rest, decrypted in-memory |
|
||||
| device_manager | Already multi-user correct (keyed by user_id), no structural change |
|
||||
|
||||
---
|
||||
|
||||
## Step 1 — WS Frame Protocol (schemas.py)
|
||||
|
||||
**Goal**: Define the v3 frame vocabulary so all subsequent steps can import it.
|
||||
|
||||
**Changes**:
|
||||
- `app/schemas.py` — Add to `WsFrameType` enum:
|
||||
- `home_request`, `floating_request`
|
||||
- `stream_start`, `stream_text`, `stream_block`, `stream_end`
|
||||
- `floating_domain`
|
||||
- `data_request`, `data_response`, `mutation`
|
||||
- Add Pydantic models:
|
||||
- `WsHomeRequest(type, message, conversation_history?)`
|
||||
- `WsFloatingRequest(type, message, scope: {type, id?})`
|
||||
- `WsStreamStart(type, request_id)`
|
||||
- `WsStreamText(type, request_id, chunk)`
|
||||
- `WsStreamBlock(type, request_id, block_type, data)`
|
||||
- `WsStreamEnd(type, request_id, mutations?)`
|
||||
- `WsFloatingDomain(type, request_id, domain)`
|
||||
- Keep all existing frame types (backward compat).
|
||||
|
||||
**Files touched**: `app/schemas.py`
|
||||
|
||||
**Test**: Unit test that validates each new model serializes/deserializes correctly.
|
||||
```
|
||||
pytest tests/test_schemas_v3.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 1 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-1: add v3 ws frame protocol (schemas.py)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 2 — Agent Streaming + Tool Result Capture (agent_registry.py, agents/)
|
||||
|
||||
**Goal**: Agents can stream LLM tokens and expose structured tool results.
|
||||
|
||||
**Changes**:
|
||||
- `app/core/agent_registry.py`:
|
||||
- Add `_tool_loop_stream()` to `ChatAgent` — same logic as `_tool_loop()` but the **final** LLM call (when no more tool calls) uses `llm.astream()` and yields tokens.
|
||||
- Add `self.tool_results: list[dict]` attribute to `ChatAgent.__init__()`.
|
||||
- In both `_tool_loop` and `_tool_loop_stream`, capture raw `execute_on_client` results when tools run (store in `self.tool_results`).
|
||||
- `app/agents/*.py` — Each agent's tools already return text summaries. No change to tools. The raw data capture happens at the `_tool_loop` level by intercepting `ToolMessage` content that comes from `execute_on_client`.
|
||||
|
||||
**Files touched**: `app/core/agent_registry.py`
|
||||
|
||||
**Test**: Unit test with mocked LLM that verifies `_tool_loop_stream()` yields tokens and `agent.tool_results` contains structured data after a tool call.
|
||||
```
|
||||
pytest tests/test_agent_streaming.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 2 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-2: add agent streaming and tool result capture (agent_registry.py)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 3 — Router Refactor (orchestrator.py)
|
||||
|
||||
**Goal**: Orchestrator returns agent name alongside execution, supports streaming.
|
||||
|
||||
**Changes**:
|
||||
- `app/core/orchestrator.py`:
|
||||
- Add `orchestrate_v3(user_id, message, context, mode)` that:
|
||||
1. Calls `classify_intent()` (unchanged) -> `agent_name`
|
||||
2. Instantiates agent via registry
|
||||
3. Returns `(agent_name, agent_instance)` — caller drives execution
|
||||
- Add `orchestrate_v3_stream(user_id, message, context)` -> `AsyncGenerator` that:
|
||||
1. Calls `classify_intent()` -> `agent_name`
|
||||
2. Calls `agent.handle_stream()` (uses `_tool_loop_stream`)
|
||||
3. Yields `(agent_name, token)` tuples — first yield includes agent name for domain detection
|
||||
- Keep `orchestrate()` and `orchestrate_stream()` unchanged (backward compat for POST /chat).
|
||||
|
||||
**Files touched**: `app/core/orchestrator.py`
|
||||
|
||||
**Test**: Unit test with mocked LLM and mocked registry that verifies `orchestrate_v3_stream` yields `(agent_name, token)` pairs.
|
||||
```
|
||||
pytest tests/test_orchestrator_v3.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 3 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-3: add router refactor with streaming support (orchestrator.py)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 4 — Output Formatting Layer (NEW: output_formatter.py)
|
||||
|
||||
**Goal**: Home and Floating responses diverge at this layer only.
|
||||
|
||||
### Block Types (from Electron app components)
|
||||
|
||||
The LLM outputs a JSON block stream. Each block has a `type` field that maps to
|
||||
an Electron renderer component. The server validates and forwards these blocks.
|
||||
|
||||
**Text block** — streamed immediately, word-by-word:
|
||||
```json
|
||||
{ "type": "text", "content": "Here's your task summary..." }
|
||||
```
|
||||
|
||||
**Chart blocks** — buffered until complete, validated, sent as `stream_block`.
|
||||
Chart types match shadcn/ui Recharts wrappers used in the Electron app:
|
||||
```json
|
||||
{ "type": "chart", "chartType": "<type>", "title": "...", "data": [...], "config": {...} }
|
||||
```
|
||||
Supported `chartType` values:
|
||||
- `area` — Area chart (shadcn AreaChart)
|
||||
- `bar` — Bar chart (shadcn BarChart)
|
||||
- `line` — Line chart (shadcn LineChart)
|
||||
- `pie` — Pie chart (shadcn PieChart)
|
||||
- `radar` — Radar chart (shadcn RadarChart)
|
||||
- `radial` — Radial/gauge chart (shadcn RadialChart)
|
||||
|
||||
`data` is an array of objects with keys matching the chart's dataKey config.
|
||||
`config` follows the shadcn ChartConfig format: `{ [dataKey]: { label, color } }`.
|
||||
|
||||
**Entity blocks** — server serializes from `agent.tool_results` (not LLM-generated data):
|
||||
```json
|
||||
{ "type": "entity_ref", "entity": "task" }
|
||||
```
|
||||
The server resolves this by looking up the structured data from the agent's
|
||||
tool call results and emitting a `stream_block` with the full entity data.
|
||||
|
||||
Supported entity types (matching Electron component types):
|
||||
- `task` — TaskRow component (`TaskItem`: id, title, status, priority, assignee, dueDate, projectId, ...)
|
||||
- `project` — Project card (id, name, clientId, status)
|
||||
- `note` — Note card (id, title, createdAt, projectId)
|
||||
- `timeline` — Timeline card (GanttTimeline: id, title, date, projectId, isAiSuggested, isApproved)
|
||||
|
||||
**Table block** — buffered, validated:
|
||||
```json
|
||||
{ "type": "table", "headers": ["Col1", "Col2"], "rows": [["val1", "val2"]] }
|
||||
```
|
||||
|
||||
**Timeline block** — buffered, validated (renders via GanttChart component):
|
||||
```json
|
||||
{ "type": "timeline", "timelines": [{ "id": "...", "title": "...", "date": 1234567890 }] }
|
||||
```
|
||||
|
||||
### Changes
|
||||
|
||||
- `app/core/output_formatter.py` (new file):
|
||||
- `HomeFormatter`:
|
||||
- Receives token stream from orchestrator
|
||||
- Accumulates tokens into a JSON-aware buffer
|
||||
- Detects block boundaries by `type` field:
|
||||
- `text` -> yields `WsStreamText` immediately (streams content word-by-word)
|
||||
- `chart` -> buffers until JSON complete, validates `chartType` against allowed set, yields `WsStreamBlock`
|
||||
- `entity_ref` -> looks up data from `agent.tool_results`, serializes full entity, yields `WsStreamBlock`
|
||||
- `table` -> buffers, validates headers/rows structure, yields `WsStreamBlock`
|
||||
- `timeline` -> buffers, validates timeline objects, yields `WsStreamBlock`
|
||||
- Invalid blocks are logged and skipped (never crash the stream)
|
||||
- `FloatingFormatter`:
|
||||
- Receives `agent_name` from orchestrator
|
||||
- Maps agent name to domain (deterministic, by code — no LLM):
|
||||
- `task_agent` -> `"tasks"`
|
||||
- `timeline_agent` -> `"timelines"`
|
||||
- `note_agent` -> `"notes"`
|
||||
- `project_agent` -> `"projects"`
|
||||
- Yields `WsFloatingDomain` immediately
|
||||
- Then yields `WsStreamText` for all tokens (text-only, no blocks)
|
||||
|
||||
**Files touched**: `app/core/output_formatter.py` (new)
|
||||
|
||||
**Test**: Unit test that feeds a mock token stream through each formatter and asserts correct frame output sequence.
|
||||
```
|
||||
pytest tests/test_output_formatter.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 4 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-4: add output formatting layer (output_formatter.py)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 5 — Unified WS Handler (device_ws.py, chat.py, main.py)
|
||||
|
||||
**Goal**: Single multiplexed WebSocket handles device frames + Home/Floating chat.
|
||||
|
||||
**Changes**:
|
||||
- `app/api/routes/device_ws.py`:
|
||||
- Extend `_message_loop` dispatch to handle `home_request` and `floating_request`:
|
||||
- On `home_request`: set `ws_context` executor, call `orchestrate_v3_stream`, pipe through `HomeFormatter`, send frames back on same socket.
|
||||
- On `floating_request`: same, but pipe through `FloatingFormatter`.
|
||||
- Wrap both in try/finally to clear `ws_context`.
|
||||
- Each request gets a `request_id` (UUID) for frame correlation.
|
||||
- Concurrent requests from same client are supported (each runs as an async task).
|
||||
- `app/api/routes/chat.py`:
|
||||
- Remove `chat_stream` WS endpoint and any related helper functions that were only used by it.
|
||||
- Keep `POST /chat` endpoint unchanged (REST fallback).
|
||||
- Clean up any unused imports.
|
||||
- `app/main.py`:
|
||||
- No change needed (device_ws router already registered).
|
||||
|
||||
**Files touched**: `app/api/routes/device_ws.py`, `app/api/routes/chat.py`, `app/main.py`
|
||||
|
||||
**Test**: Integration test with a WebSocket test client that:
|
||||
1. Connects to `/api/v1/ws/device`
|
||||
2. Sends `device_hello`
|
||||
3. Sends `home_request` -> receives `stream_start`, `stream_text`*, `stream_end`
|
||||
4. Sends `floating_request` -> receives `floating_domain`, `stream_text`*, `stream_end`
|
||||
5. Verifies `tool_call`/`tool_result` round-trip still works during chat
|
||||
```
|
||||
pytest tests/test_ws_unified.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 5 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-5: unify ws handler (device_ws.py, chat.py)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 6 — Memory Models + Migration (models.py, alembic)
|
||||
|
||||
**Goal**: Database tables for 4-tier memory, with per-user encryption key.
|
||||
|
||||
**Changes**:
|
||||
- `app/models.py`:
|
||||
- Add `encryption_key` column to `User` model (Fernet key, generated on registration).
|
||||
- Add `MemoryCore` model: `id, user_id, key, value_encrypted, updated_at`
|
||||
- Add `MemoryAssociative` model: `id, user_id, content_encrypted, embedding (Vector(1536)), entity_type, entity_id, updated_at`
|
||||
- Add `MemoryEpisodic` model: `id, user_id, summary_encrypted, session_id, created_at`
|
||||
- Add `MemoryProactive` model: `id, user_id, pattern_encrypted, confidence, source, created_at`
|
||||
- `alembic/versions/` — New migration adding the 4 memory tables + user encryption_key column.
|
||||
- `app/api/routes/auth.py` — On user registration, generate and store a Fernet key.
|
||||
|
||||
**Files touched**: `app/models.py`, `alembic/versions/xxx_add_memory_tables.py`, `app/api/routes/auth.py`
|
||||
|
||||
**Test**: Run migration up/down, verify tables exist with correct columns.
|
||||
```
|
||||
alembic upgrade head && alembic downgrade -1 && alembic upgrade head
|
||||
pytest tests/test_memory_models.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 6 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-6: add memory models and migration (models.py, alembic)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 7 — Memory Middleware (NEW: memory_middleware.py)
|
||||
|
||||
**Goal**: Enrich every Router call with memory context, store interactions after.
|
||||
|
||||
**Changes**:
|
||||
- `app/core/memory_middleware.py` (new file):
|
||||
- `MemoryMiddleware` class with:
|
||||
- `enrich_context(user_id, message) -> dict` (pre-LLM):
|
||||
1. Load core memory (user prefs) — always injected
|
||||
2. Embed `message`, search `MemoryAssociative` via pgvector — top-k relevant
|
||||
3. Fetch recent `MemoryEpisodic` entries — last N sessions
|
||||
4. Fetch active `MemoryProactive` patterns — above confidence threshold
|
||||
5. Return merged context dict
|
||||
- `store_episode(user_id, session_id, message, response)` (post-LLM):
|
||||
1. Summarize interaction (short LLM call or heuristic)
|
||||
2. Encrypt and store in `MemoryEpisodic`
|
||||
3. Embed interaction, encrypt and upsert in `MemoryAssociative`
|
||||
- `update_core(user_id, key, value)` — explicit preference update
|
||||
- All read/write operations encrypt/decrypt using the user's Fernet key from `User.encryption_key`
|
||||
- `app/api/routes/device_ws.py` — Update `home_request` and `floating_request` handlers:
|
||||
- Before orchestrator: `enriched = await memory.enrich_context(user_id, message)`
|
||||
- After response complete: `await memory.store_episode(user_id, ...)`
|
||||
|
||||
**Files touched**: `app/core/memory_middleware.py` (new), `app/api/routes/device_ws.py`
|
||||
|
||||
**Test**: Unit test with seeded memory rows that verifies:
|
||||
1. `enrich_context` returns core prefs + associative matches + episodic summaries
|
||||
2. `store_episode` creates encrypted rows that can be decrypted with the user's key
|
||||
3. End-to-end WS test: send `home_request`, verify memory enrichment is passed to orchestrator
|
||||
```
|
||||
pytest tests/test_memory_middleware.py
|
||||
```
|
||||
|
||||
**Status**:
|
||||
- [x] Step 7 complete
|
||||
|
||||
**Commit**: After tests pass, commit with:
|
||||
```
|
||||
git commit -m "step-7: add memory middleware (memory_middleware.py, device_ws.py)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
| Step | Component | Effort | Depends On |
|
||||
|------|-----------|--------|------------|
|
||||
| 1 | WS Frame Protocol | Low | — |
|
||||
| 2 | Agent Streaming | Medium | Step 1 |
|
||||
| 3 | Router Refactor | Medium | Step 2 |
|
||||
| 4 | Output Formatter | High | Steps 1, 3 |
|
||||
| 5 | Unified WS Handler | High | Steps 1–4 |
|
||||
| 6 | Memory Models | Medium | — |
|
||||
| 7 | Memory Middleware | High | Steps 5, 6 |
|
||||
|
||||
Steps 1–5 form the streaming pipeline. Steps 6–7 form the memory system.
|
||||
Step 6 can run in parallel with Steps 2–4 (no dependencies).
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Expose tool modules used by deep orchestrator-worker graphs."""
|
||||
"""Agent tool modules — imported by deep_agent.py to build sub-agent graphs."""
|
||||
|
||||
from app.agents import timeline_agent, note_agent, project_agent, task_agent
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Note agent — Markdown note management (list, get, create, update, delete)."""
|
||||
"""Note agent — tool definitions for Markdown note CRUD."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -9,20 +9,6 @@ from langchain_core.tools import tool
|
||||
from app.core.llm import embed
|
||||
from app.core.ws_context import execute_on_client
|
||||
|
||||
NOTE_SYSTEM_PROMPT = (
|
||||
"You are a note-taking assistant. You help users create, retrieve, update,\n"
|
||||
"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)."
|
||||
)
|
||||
|
||||
|
||||
@tool
|
||||
async def list_notes(project_id: str = "") -> str:
|
||||
@@ -119,10 +105,4 @@ async def delete_note(note_id: str) -> str:
|
||||
return f"Note {note_id} deleted."
|
||||
|
||||
|
||||
NOTE_TOOLS: list[Any] = [
|
||||
list_notes,
|
||||
get_note,
|
||||
create_note,
|
||||
update_note,
|
||||
delete_note,
|
||||
]
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Project agent — full lifecycle management (list, get, create, update, archive, delete)."""
|
||||
"""Project agent — tool definitions for project lifecycle CRUD."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -8,22 +8,6 @@ from langchain_core.tools import tool
|
||||
|
||||
from app.core.ws_context import execute_on_client
|
||||
|
||||
PROJECT_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
|
||||
async def list_projects(
|
||||
@@ -133,11 +117,4 @@ async def delete_project(project_id: str) -> str:
|
||||
return f"Project {project_id} permanently deleted."
|
||||
|
||||
|
||||
PROJECT_TOOLS: list[Any] = [
|
||||
list_projects,
|
||||
list_all_projects,
|
||||
get_project,
|
||||
create_project,
|
||||
update_project,
|
||||
delete_project,
|
||||
]
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Task agent — full CRUD for tasks and task comments."""
|
||||
"""Task agent — tool definitions for task and task comment CRUD."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -9,23 +9,6 @@ from langchain_core.tools import tool
|
||||
|
||||
from app.core.ws_context import execute_on_client
|
||||
|
||||
TASK_SYSTEM_PROMPT = (
|
||||
"You are a task management assistant for a project workspace.\n"
|
||||
"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."
|
||||
)
|
||||
|
||||
|
||||
# ── Task tools ────────────────────────────────────────────────────────
|
||||
|
||||
@@ -216,16 +199,4 @@ async def delete_task_comment(comment_id: str) -> str:
|
||||
return f"Comment {comment_id} deleted."
|
||||
|
||||
|
||||
# ── Agent ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
TASK_TOOLS: list[Any] = [
|
||||
list_tasks,
|
||||
create_task,
|
||||
update_task,
|
||||
delete_task,
|
||||
list_tasks_due_today,
|
||||
list_task_comments,
|
||||
add_task_comment,
|
||||
delete_task_comment,
|
||||
]
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Timeline agent — project milestone management (list, create, update, delete)."""
|
||||
"""Timeline agent — tool definitions for project milestone CRUD."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -8,19 +8,6 @@ from langchain_core.tools import tool
|
||||
|
||||
from app.core.ws_context import execute_on_client
|
||||
|
||||
TIMELINE_SYSTEM_PROMPT = (
|
||||
"You are a project timeline assistant. Timelines 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 timeline, 0 otherwise\n"
|
||||
" - is_approved: 0 until the user explicitly confirms; then 1\n"
|
||||
" - For update_timeline, use -1 for integer fields you do not want to change\n"
|
||||
" - Listing without a project_id returns all timelines across projects\n"
|
||||
" - Always echo the title and formatted date in your confirmation."
|
||||
)
|
||||
|
||||
|
||||
@tool
|
||||
async def list_timelines(project_id: str = "") -> str:
|
||||
@@ -102,9 +89,4 @@ async def delete_timeline(timeline_id: str) -> str:
|
||||
return f"Timeline {timeline_id} deleted."
|
||||
|
||||
|
||||
TIMELINE_TOOLS: list[Any] = [
|
||||
list_timelines,
|
||||
create_timeline,
|
||||
update_timeline,
|
||||
delete_timeline,
|
||||
]
|
||||
|
||||
|
||||
@@ -10,7 +10,9 @@ from fastapi.responses import JSONResponse
|
||||
|
||||
from app.api.deps import get_current_user
|
||||
from app.core.deep_agent import run_home
|
||||
from app.schemas import ChatRequest, UserProfile
|
||||
from app.core.memory_middleware import MemoryMiddleware
|
||||
from app.db import async_session
|
||||
from app.schemas import ChatRequest, ChatResponse, UserProfile
|
||||
|
||||
router = APIRouter(prefix="/chat", tags=["chat"])
|
||||
|
||||
@@ -20,10 +22,21 @@ async def chat(
|
||||
body: ChatRequest,
|
||||
current_user: UserProfile = Depends(get_current_user),
|
||||
) -> JSONResponse:
|
||||
"""REST fallback for home chat when websocket streaming is unavailable."""
|
||||
response = await run_home(
|
||||
"""Route a chat message through the Home deep agent (non-streaming)."""
|
||||
async with async_session() as db:
|
||||
memory = MemoryMiddleware(db)
|
||||
memory_context = await memory.enrich_context(current_user.id, body.message)
|
||||
|
||||
context = {
|
||||
**body.context.model_dump(),
|
||||
**memory_context,
|
||||
}
|
||||
|
||||
response_text = await run_home(
|
||||
user_id=current_user.id,
|
||||
message=body.message,
|
||||
context=body.context.model_dump(),
|
||||
context=context,
|
||||
db_session_factory=async_session,
|
||||
)
|
||||
return JSONResponse(content={"response": response})
|
||||
result = ChatResponse(response=response_text)
|
||||
return JSONResponse(content=result.model_dump())
|
||||
|
||||
@@ -41,10 +41,10 @@ from sqlalchemy import update
|
||||
|
||||
from app.config.settings import settings
|
||||
from app.core.agent_runner import trigger_pending_runs
|
||||
from app.core.deep_agent import run_floating_stream, run_home_stream
|
||||
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.deep_agent import run_home_stream, run_floating_stream
|
||||
from app.core.output_formatter import HomeFormatter, FloatingFormatter
|
||||
from app.core.ws_context import clear_client_executor, set_client_executor
|
||||
from app.db import async_session
|
||||
from app.models import AgentRunLog
|
||||
@@ -200,13 +200,35 @@ async def _message_loop(websocket: WebSocket, user_id: str) -> None:
|
||||
|
||||
# ── v3 Chat Handlers ──────────────────────────────────────────────────
|
||||
|
||||
_WS_TOOL_CALL_TIMEOUT = 30 # seconds to wait for Electron tool_result
|
||||
|
||||
|
||||
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
|
||||
call_id = payload["id"]
|
||||
logger.info("ws_executor: sending tool_call id=%s action=%s", call_id, payload.get("action"))
|
||||
await websocket.send_text(json.dumps(payload))
|
||||
future = device_manager.create_pending_call(user_id, payload["id"])
|
||||
return await future
|
||||
future = device_manager.create_pending_call(user_id, call_id)
|
||||
try:
|
||||
result = await asyncio.wait_for(future, timeout=_WS_TOOL_CALL_TIMEOUT)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
"ws_executor: timeout waiting for tool_result id=%s action=%s user=%s",
|
||||
call_id, payload.get("action"), user_id,
|
||||
)
|
||||
# Clean up the pending future so it doesn't leak
|
||||
conn = device_manager._connections.get(user_id)
|
||||
if conn:
|
||||
conn.pending_calls.pop(call_id, None)
|
||||
return {"error": f"Tool call timed out after {_WS_TOOL_CALL_TIMEOUT}s", "rows": []}
|
||||
logger.info("ws_executor: tool_result id=%s result_type=%s result_keys=%s",
|
||||
call_id, type(result).__name__,
|
||||
list(result.keys()) if isinstance(result, dict) else "N/A")
|
||||
if result is None:
|
||||
logger.error("ws_executor: future resolved to None for call_id=%s user=%s", call_id, user_id)
|
||||
return result
|
||||
return _executor
|
||||
|
||||
|
||||
@@ -234,11 +256,12 @@ async def _handle_home_request(
|
||||
set_client_executor(executor)
|
||||
response_chunks: list[str] = []
|
||||
try:
|
||||
event_stream = run_home_stream(user_id, message, context)
|
||||
formatter = StreamFormatter(request_id=request_id)
|
||||
event_stream = run_home_stream(
|
||||
user_id, message, context, db_session_factory=async_session
|
||||
)
|
||||
formatter = HomeFormatter(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:
|
||||
@@ -279,8 +302,11 @@ async def _handle_floating_request(
|
||||
set_client_executor(executor)
|
||||
response_chunks: list[str] = []
|
||||
try:
|
||||
event_stream = run_floating_stream(user_id, message, context)
|
||||
formatter = StreamFormatter(request_id=request_id)
|
||||
event_stream = run_floating_stream(
|
||||
user_id, message, context, scope=scope,
|
||||
db_session_factory=async_session,
|
||||
)
|
||||
formatter = FloatingFormatter(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]
|
||||
|
||||
@@ -1,30 +0,0 @@
|
||||
"""Minimal agent base types retained for compatibility with batch runners."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
|
||||
class BaseAgent(ABC):
|
||||
"""Common base for non-chat agents still using the old base contract."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str = "",
|
||||
shared_memory: dict[str, Any] | None = None,
|
||||
vector_store_context: list[str] | None = None,
|
||||
) -> None:
|
||||
self.user_id = user_id
|
||||
self.shared_memory: dict[str, Any] = shared_memory or {}
|
||||
self.vector_store_context: list[str] = vector_store_context or []
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> str: ...
|
||||
|
||||
@abstractmethod
|
||||
def get_description(self) -> str: ...
|
||||
|
||||
@property
|
||||
def skills(self) -> list[str]:
|
||||
return []
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Agent run orchestrator.
|
||||
"""Agent run manager.
|
||||
|
||||
Drives two agent types:
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
"""LLM factory — centralised model instantiation via LiteLLM.
|
||||
|
||||
Every agent and the orchestrator call ``get_llm()`` or ``get_router_llm()``
|
||||
Every agent and the deep-agent supervisors call ``get_llm()`` or ``get_router_llm()``
|
||||
instead of directly constructing a provider-specific class. The model string
|
||||
follows the `LiteLLM model naming convention
|
||||
<https://docs.litellm.ai/docs/providers>`_:
|
||||
|
||||
@@ -43,7 +43,7 @@ _PROACTIVE_CONFIDENCE_THRESHOLD = 0.6
|
||||
|
||||
|
||||
class MemoryMiddleware:
|
||||
"""Enrich orchestrator context with memory and persist interactions after."""
|
||||
"""Enrich agent context with memory and persist interactions after."""
|
||||
|
||||
def __init__(self, db: AsyncSession) -> None:
|
||||
self._db = db
|
||||
@@ -51,7 +51,7 @@ class MemoryMiddleware:
|
||||
# ── Public API ────────────────────────────────────────────────────────────
|
||||
|
||||
async def enrich_context(self, user_id: str, message: str) -> dict[str, Any]:
|
||||
"""Build memory context dict to inject into the orchestrator before LLM call.
|
||||
"""Build memory context dict to inject into the agent before LLM call.
|
||||
|
||||
Returns a dict with keys:
|
||||
core_memory — {key: plaintext_value, ...}
|
||||
|
||||
@@ -1,43 +1,141 @@
|
||||
"""Output formatter for deep-agent stream events."""
|
||||
"""Output Formatter — transforms deep-agent event streams into WS frame sequences.
|
||||
|
||||
Consumes ``(event_type, data)`` tuples yielded by ``deep_agent.run_*_stream()``:
|
||||
* ``("token", str)`` — supervisor text token
|
||||
* ``("tool_end", dict)`` — sub-agent finished: ``{name, result}``
|
||||
* ``("mutations", list)`` — collected CRUD mutations for ``stream_end``
|
||||
|
||||
HomeFormatter:
|
||||
* Streams text tokens as-is → emits ``WsStreamText``
|
||||
(text may contain inline ``<type>[id,...]</type>`` entity tags
|
||||
for the frontend to parse and render as interactive components)
|
||||
* Attaches mutations → injects into ``WsStreamEnd``
|
||||
|
||||
FloatingFormatter:
|
||||
* Sniffs first ``tool_end`` name → emits ``WsFloatingDomain``
|
||||
* Streams text tokens → emits ``WsStreamText``
|
||||
* Attaches mutations → injects into ``WsStreamEnd``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
from app.schemas import WsFloatingDomain, WsStreamEnd, WsStreamStart, WsStreamText
|
||||
from app.schemas import (
|
||||
WsFloatingDomain,
|
||||
WsStreamEnd,
|
||||
WsStreamStart,
|
||||
WsStreamText,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Map sub-agent tool name → floating domain / entity type
|
||||
_AGENT_DOMAIN: dict[str, str] = {
|
||||
"task_agent": "tasks",
|
||||
"timeline_agent": "timelines",
|
||||
"note_agent": "notes",
|
||||
"project_agent": "projects",
|
||||
}
|
||||
|
||||
WsFrame = WsStreamStart | WsStreamText | WsStreamEnd | WsFloatingDomain
|
||||
|
||||
|
||||
class StreamFormatter:
|
||||
"""Convert `(event_type, data)` stream events into websocket frame models."""
|
||||
class HomeFormatter:
|
||||
"""Consumes a deep-agent event stream and yields WS frames for the Home view.
|
||||
|
||||
Text tokens are forwarded as-is via ``WsStreamText``. The supervisor
|
||||
embeds ``<type>[id1,id2]</type>`` entity tags inline — the frontend
|
||||
is responsible for parsing those and rendering interactive components.
|
||||
Mutations are attached to ``WsStreamEnd``.
|
||||
"""
|
||||
|
||||
def __init__(self, request_id: str) -> None:
|
||||
self.request_id = request_id
|
||||
self._mutations: list[dict] = []
|
||||
|
||||
async def format(
|
||||
self,
|
||||
event_stream: AsyncGenerator[tuple[str, Any], None],
|
||||
) -> AsyncGenerator[WsFrame, None]:
|
||||
started = False
|
||||
yield WsStreamStart(request_id=self.request_id)
|
||||
|
||||
async for event_type, data in event_stream:
|
||||
if event_type == "floating_domain":
|
||||
yield WsFloatingDomain(request_id=self.request_id, domain=str(data))
|
||||
continue
|
||||
if event_type == "token":
|
||||
if data:
|
||||
yield WsStreamText(request_id=self.request_id, chunk=data)
|
||||
|
||||
if event_type != "token":
|
||||
continue
|
||||
elif event_type == "mutations":
|
||||
self._mutations = data or []
|
||||
|
||||
if not started:
|
||||
yield WsStreamEnd(
|
||||
request_id=self.request_id,
|
||||
mutations=[
|
||||
{"action": m["action"], "table": m["table"], "data": m["data"]}
|
||||
for m in self._mutations
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class FloatingFormatter:
|
||||
"""Consumes a deep-agent event stream and yields WS frames for the Floating view.
|
||||
|
||||
Sniffs the first ``tool_end`` event name to derive the domain (e.g.
|
||||
``task_agent`` → ``"tasks"``), then streams text tokens as plain
|
||||
``WsStreamText``. No block parsing for floating context.
|
||||
"""
|
||||
|
||||
def __init__(self, request_id: str) -> None:
|
||||
self.request_id = request_id
|
||||
self._mutations: list[dict] = []
|
||||
|
||||
async def format(
|
||||
self,
|
||||
event_stream: AsyncGenerator[tuple[str, Any], None],
|
||||
) -> AsyncGenerator[WsFrame, None]:
|
||||
domain_sent = False
|
||||
|
||||
async for event_type, data in event_stream:
|
||||
if event_type == "tool_end" and not domain_sent:
|
||||
# Sniff domain from the first sub-agent that completes
|
||||
name = data.get("name", "")
|
||||
domain = _AGENT_DOMAIN.get(name, "tasks")
|
||||
yield WsFloatingDomain(
|
||||
request_id=self.request_id,
|
||||
domain=domain, # type: ignore[arg-type]
|
||||
)
|
||||
yield WsStreamStart(request_id=self.request_id)
|
||||
started = True
|
||||
domain_sent = True
|
||||
|
||||
text = str(data or "")
|
||||
if text:
|
||||
yield WsStreamText(request_id=self.request_id, chunk=text)
|
||||
elif event_type == "token":
|
||||
if not domain_sent:
|
||||
# First token arrived before any tool_end — default domain
|
||||
yield WsFloatingDomain(
|
||||
request_id=self.request_id,
|
||||
domain="tasks", # type: ignore[arg-type]
|
||||
)
|
||||
yield WsStreamStart(request_id=self.request_id)
|
||||
domain_sent = True
|
||||
if data:
|
||||
yield WsStreamText(request_id=self.request_id, chunk=data)
|
||||
|
||||
if not started:
|
||||
elif event_type == "mutations":
|
||||
self._mutations = data or []
|
||||
|
||||
# If no events triggered domain_sent (edge case), still emit structure
|
||||
if not domain_sent:
|
||||
yield WsFloatingDomain(
|
||||
request_id=self.request_id,
|
||||
domain="tasks", # type: ignore[arg-type]
|
||||
)
|
||||
yield WsStreamStart(request_id=self.request_id)
|
||||
yield WsStreamEnd(request_id=self.request_id)
|
||||
|
||||
yield WsStreamEnd(
|
||||
request_id=self.request_id,
|
||||
mutations=[
|
||||
{"action": m["action"], "table": m["table"], "data": m["data"]}
|
||||
for m in self._mutations
|
||||
],
|
||||
)
|
||||
|
||||
@@ -7,18 +7,21 @@ The callback sends a `tool_call` WS frame and awaits the `tool_result`.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from contextvars import ContextVar
|
||||
from typing import Any, Callable, Coroutine
|
||||
from uuid import uuid4
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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 deep agent runs; cleared after.
|
||||
_client_executor: ContextVar[Callable[[dict], Coroutine[Any, Any, dict]]] = ContextVar(
|
||||
"_client_executor"
|
||||
)
|
||||
|
||||
# Optional collector that captures raw execute_on_client results.
|
||||
# Set by _tool_loop / _tool_loop_stream to populate ChatAgent.tool_results.
|
||||
# Set by the deep agent tool loop to capture CRUD mutations.
|
||||
_tool_result_collector: ContextVar[list[dict] | None] = ContextVar(
|
||||
"_tool_result_collector", default=None
|
||||
)
|
||||
@@ -81,12 +84,17 @@ async def execute_on_client(
|
||||
if limit is not None:
|
||||
payload["limit"] = limit
|
||||
|
||||
logger.info("execute_on_client: sending payload action=%s table=%s id=%s", action, table, payload["id"])
|
||||
result = await callback(payload)
|
||||
if result is None:
|
||||
logger.error("execute_on_client: callback returned None for action=%s table=%s id=%s", action, table, payload["id"])
|
||||
else:
|
||||
logger.info("execute_on_client: got result type=%s keys=%s", type(result).__name__, list(result.keys()) if isinstance(result, dict) else "N/A")
|
||||
collector = _tool_result_collector.get(None)
|
||||
if collector is not None:
|
||||
if collector is not None and action in ("insert", "update", "delete"):
|
||||
collector.append({
|
||||
"action": action,
|
||||
"table": table,
|
||||
"data": result,
|
||||
"data": data or {},
|
||||
})
|
||||
return result
|
||||
|
||||
@@ -18,9 +18,7 @@ from app.config.settings import settings
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
# Startup: ensure agent tool modules are loaded.
|
||||
import app.agents # noqa: F401
|
||||
|
||||
# Startup: initialise DB connection pool
|
||||
yield
|
||||
|
||||
# Shutdown: dispose SQLAlchemy connection pool
|
||||
|
||||
@@ -279,6 +279,7 @@ class WsStreamEnd(BaseModel):
|
||||
|
||||
type: Literal[WsFrameType.stream_end] = WsFrameType.stream_end
|
||||
request_id: str
|
||||
mutations: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
|
||||
class WsFloatingDomain(BaseModel):
|
||||
|
||||
@@ -4,8 +4,9 @@ gunicorn>=22.0.0
|
||||
langchain>=0.3.0
|
||||
langchain-openai>=0.3.0
|
||||
langchain-litellm>=0.1.0
|
||||
langgraph>=0.3.0
|
||||
deepagents>=0.4.10
|
||||
litellm>=1.50.0
|
||||
langgraph>=0.4.0
|
||||
pydantic>=2.10.0
|
||||
pydantic-settings>=2.7.0
|
||||
python-jose[cryptography]>=3.3.0
|
||||
|
||||
@@ -250,10 +250,11 @@ def test_home_request_calls_memory_middleware(client):
|
||||
token = make_jwt("power", user_id=USER_ID)
|
||||
session_id = str(uuid.uuid4())
|
||||
|
||||
async def _mock_stream(user_id, message, context):
|
||||
async def _mock_stream(user_id, message, context, db_session_factory=None):
|
||||
# Verify memory context was injected
|
||||
assert context.get("core_memory") == {"tz": "UTC"}
|
||||
yield "token", "Done"
|
||||
yield ("token", "Done")
|
||||
yield ("mutations", [])
|
||||
|
||||
with (
|
||||
patch("app.api.routes.device_ws.MemoryMiddleware", _MockMiddleware),
|
||||
|
||||
@@ -1,75 +1,214 @@
|
||||
"""Tests for app.core.output_formatter.StreamFormatter."""
|
||||
"""Tests for app.core.output_formatter — HomeFormatter and FloatingFormatter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.output_formatter import StreamFormatter
|
||||
from app.schemas import WsFloatingDomain, WsStreamEnd, WsStreamStart, WsStreamText
|
||||
from app.core.output_formatter import HomeFormatter, FloatingFormatter
|
||||
from app.schemas import (
|
||||
WsFloatingDomain,
|
||||
WsStreamEnd,
|
||||
WsStreamStart,
|
||||
WsStreamText,
|
||||
)
|
||||
|
||||
|
||||
# ── helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
async def _stream(*events: tuple[str, object]):
|
||||
"""Async generator that yields (event_type, data) tuples."""
|
||||
for event in events:
|
||||
yield event
|
||||
|
||||
|
||||
async def _collect(formatter: StreamFormatter, event_stream):
|
||||
async def collect(formatter, event_stream):
|
||||
frames = []
|
||||
async for frame in formatter.format(event_stream):
|
||||
frames.append(frame)
|
||||
return frames
|
||||
|
||||
|
||||
# ── HomeFormatter ─────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_formatter_text_stream() -> None:
|
||||
formatter = StreamFormatter(request_id="req-1")
|
||||
frames = await _collect(
|
||||
formatter,
|
||||
_stream(("token", "Hello"), ("token", " world")),
|
||||
)
|
||||
async def test_home_formatter_plain_text():
|
||||
req_id = "req-1"
|
||||
events = [
|
||||
("token", "Hello world"),
|
||||
("mutations", []),
|
||||
]
|
||||
formatter = HomeFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
assert isinstance(frames[0], WsStreamStart)
|
||||
assert isinstance(frames[1], WsStreamText)
|
||||
assert frames[1].chunk == "Hello"
|
||||
assert isinstance(frames[2], WsStreamText)
|
||||
assert frames[2].chunk == " world"
|
||||
assert frames[0].request_id == req_id
|
||||
text_frames = [f for f in frames if isinstance(f, WsStreamText)]
|
||||
assert any("Hello world" in f.chunk for f in text_frames)
|
||||
assert isinstance(frames[-1], WsStreamEnd)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_formatter_floating_domain_first() -> None:
|
||||
formatter = StreamFormatter(request_id="req-2")
|
||||
frames = await _collect(
|
||||
formatter,
|
||||
_stream(("floating_domain", "notes"), ("token", "Summary")),
|
||||
)
|
||||
async def test_home_formatter_entity_tags_passed_through():
|
||||
"""Entity tags are streamed as-is — the frontend parses them."""
|
||||
req_id = "req-2"
|
||||
events = [
|
||||
("token", "Here is your project:\n<project>[abc-123]</project>\nAll good."),
|
||||
("mutations", []),
|
||||
]
|
||||
formatter = HomeFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
text = "".join(f.chunk for f in frames if isinstance(f, WsStreamText))
|
||||
assert "<project>[abc-123]</project>" in text
|
||||
assert "Here is your project:" in text
|
||||
assert "All good." in text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_home_formatter_multiple_tags_passed_through():
|
||||
req_id = "req-3"
|
||||
events = [
|
||||
("token", "<project>[p1]</project>\nText\n<task>[t1,t2]</task>"),
|
||||
("mutations", []),
|
||||
]
|
||||
formatter = HomeFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
text = "".join(f.chunk for f in frames if isinstance(f, WsStreamText))
|
||||
assert "<project>[p1]</project>" in text
|
||||
assert "<task>[t1,t2]</task>" in text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_home_formatter_tool_end_ignored():
|
||||
"""tool_end events are silently ignored by HomeFormatter."""
|
||||
req_id = "req-4"
|
||||
events = [
|
||||
("tool_end", {"name": "task_agent", "result": "3 tasks"}),
|
||||
("token", "No tags here."),
|
||||
("mutations", []),
|
||||
]
|
||||
formatter = HomeFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
text = "".join(f.chunk for f in frames if isinstance(f, WsStreamText))
|
||||
assert text == "No tags here."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_home_formatter_mutations_in_stream_end():
|
||||
req_id = "req-5"
|
||||
muts = [{"action": "insert", "table": "tasks", "data": {"id": "t1"}}]
|
||||
events = [
|
||||
("token", "Done"),
|
||||
("mutations", muts),
|
||||
]
|
||||
formatter = HomeFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
end_frame = frames[-1]
|
||||
assert isinstance(end_frame, WsStreamEnd)
|
||||
assert len(end_frame.mutations) == 1
|
||||
assert end_frame.mutations[0]["action"] == "insert"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_home_formatter_frame_order():
|
||||
"""stream_start is first, stream_end is last."""
|
||||
req_id = "req-6"
|
||||
formatter = HomeFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(("token", "Hi"), ("mutations", [])))
|
||||
assert isinstance(frames[0], WsStreamStart)
|
||||
assert isinstance(frames[-1], WsStreamEnd)
|
||||
|
||||
|
||||
# ── FloatingFormatter ─────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_floating_formatter_domain_from_tool_end():
|
||||
req_id = "pop-1"
|
||||
formatter = FloatingFormatter(request_id=req_id)
|
||||
events = [
|
||||
("tool_end", {"name": "task_agent", "result": "ok"}),
|
||||
("token", "Hello"),
|
||||
("mutations", []),
|
||||
]
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
assert isinstance(frames[0], WsFloatingDomain)
|
||||
assert frames[0].domain == "notes"
|
||||
assert isinstance(frames[1], WsStreamStart)
|
||||
assert isinstance(frames[2], WsStreamText)
|
||||
assert frames[2].chunk == "Summary"
|
||||
assert frames[0].domain == "tasks"
|
||||
assert frames[0].request_id == req_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_floating_formatter_text_only():
|
||||
req_id = "pop-2"
|
||||
formatter = FloatingFormatter(request_id=req_id)
|
||||
events = [
|
||||
("tool_end", {"name": "timeline_agent", "result": "done"}),
|
||||
("token", "Summary"),
|
||||
("mutations", []),
|
||||
]
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
assert isinstance(frames[0], WsFloatingDomain)
|
||||
assert frames[0].domain == "timelines"
|
||||
text_frames = [f for f in frames if isinstance(f, WsStreamText)]
|
||||
assert len(text_frames) == 1
|
||||
assert text_frames[0].chunk == "Summary"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_floating_formatter_no_entity_tags():
|
||||
"""FloatingFormatter never emits entity tag blocks."""
|
||||
req_id = "pop-3"
|
||||
formatter = FloatingFormatter(request_id=req_id)
|
||||
events = [
|
||||
("tool_end", {"name": "note_agent", "result": "data"}),
|
||||
("token", "some text"),
|
||||
("mutations", []),
|
||||
]
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
# Only expected frame types
|
||||
for f in frames:
|
||||
assert isinstance(f, (WsFloatingDomain, WsStreamStart, WsStreamText, WsStreamEnd))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_floating_formatter_end_frame():
|
||||
req_id = "pop-4"
|
||||
formatter = FloatingFormatter(request_id=req_id)
|
||||
events = [
|
||||
("tool_end", {"name": "project_agent", "result": "ok"}),
|
||||
("token", "Done"),
|
||||
("mutations", []),
|
||||
]
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
assert isinstance(frames[-1], WsStreamEnd)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_formatter_ignores_unknown_events() -> None:
|
||||
formatter = StreamFormatter(request_id="req-3")
|
||||
frames = await _collect(
|
||||
formatter,
|
||||
_stream(("tool_end", {"name": "x"}), ("token", "ok")),
|
||||
)
|
||||
|
||||
text_frames = [f for f in frames if isinstance(f, WsStreamText)]
|
||||
assert len(text_frames) == 1
|
||||
assert text_frames[0].chunk == "ok"
|
||||
async def test_floating_formatter_default_domain_on_early_token():
|
||||
"""When the first event is a token (no tool_end yet), default to 'tasks'."""
|
||||
req_id = "pop-5"
|
||||
formatter = FloatingFormatter(request_id=req_id)
|
||||
events = [("token", "hi"), ("mutations", [])]
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
assert isinstance(frames[0], WsFloatingDomain)
|
||||
assert frames[0].domain == "tasks"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_formatter_empty_stream_still_brackets() -> None:
|
||||
formatter = StreamFormatter(request_id="req-4")
|
||||
frames = await _collect(formatter, _stream())
|
||||
async def test_floating_formatter_mutations_in_stream_end():
|
||||
req_id = "pop-6"
|
||||
muts = [{"action": "update", "table": "tasks", "data": {"id": "t2"}}]
|
||||
events = [
|
||||
("token", "Updated"),
|
||||
("mutations", muts),
|
||||
]
|
||||
formatter = FloatingFormatter(request_id=req_id)
|
||||
frames = await collect(formatter, _stream(*events))
|
||||
|
||||
assert len(frames) == 2
|
||||
assert isinstance(frames[0], WsStreamStart)
|
||||
assert isinstance(frames[1], WsStreamEnd)
|
||||
end_frame = frames[-1]
|
||||
assert isinstance(end_frame, WsStreamEnd)
|
||||
assert len(end_frame.mutations) == 1
|
||||
|
||||
@@ -88,7 +88,7 @@ class TestPluginRegistry:
|
||||
async def test_list_filter_by_query(
|
||||
self, reg: PluginRegistry, db_session: AsyncSession, seed_plugins: list[Plugin]
|
||||
) -> None:
|
||||
result = await reg.list_plugins(db_session, query="time")
|
||||
result = await reg.list_plugins(db_session, query="time tracker")
|
||||
assert result.total == 1
|
||||
assert result.plugins[0].id == "plugin-time-tracker"
|
||||
|
||||
|
||||
@@ -178,15 +178,23 @@ def test_stream_text_deserializes():
|
||||
def test_stream_end_defaults():
|
||||
frame = WsStreamEnd(request_id="r1")
|
||||
assert frame.type == WsFrameType.stream_end
|
||||
assert frame.mutations == []
|
||||
|
||||
|
||||
def test_stream_end_with_mutations():
|
||||
mutations = [{"action": "create", "table": "tasks", "data": {"title": "New task"}}]
|
||||
frame = WsStreamEnd(request_id="r1", mutations=mutations)
|
||||
assert len(frame.mutations) == 1
|
||||
assert frame.mutations[0]["action"] == "create"
|
||||
|
||||
|
||||
def test_stream_end_serializes():
|
||||
data = WsStreamEnd(request_id="r2").model_dump()
|
||||
assert data == {"type": "stream_end", "request_id": "r2"}
|
||||
assert data == {"type": "stream_end", "request_id": "r2", "mutations": []}
|
||||
|
||||
|
||||
def test_stream_end_deserializes():
|
||||
raw = {"type": "stream_end", "request_id": "r3"}
|
||||
raw = {"type": "stream_end", "request_id": "r3", "mutations": []}
|
||||
frame = WsStreamEnd.model_validate(raw)
|
||||
assert frame.request_id == "r3"
|
||||
|
||||
|
||||
@@ -45,13 +45,15 @@ def _recv_until_end(ws, max_frames: int = 20) -> list[dict]:
|
||||
return frames
|
||||
|
||||
|
||||
async def _mock_home_stream(user_id, message, context):
|
||||
yield "token", "Hello"
|
||||
async def _mock_home_stream(user_id, message, context, db_session_factory=None):
|
||||
yield "token", "Here are your tasks:\n<task>[t1,t2]</task>"
|
||||
yield "mutations", []
|
||||
|
||||
|
||||
async def _mock_floating_stream(user_id, message, context):
|
||||
yield "floating_domain", "tasks"
|
||||
async def _mock_floating_stream(user_id, message, context, scope=None, db_session_factory=None):
|
||||
yield "tool_end", {"name": "task_agent", "result": "ok"}
|
||||
yield "token", "Here is a summary"
|
||||
yield "mutations", []
|
||||
|
||||
|
||||
# ── tests ─────────────────────────────────────────────────────────────────────
|
||||
@@ -111,8 +113,9 @@ def test_home_request_request_id_propagated(client):
|
||||
token = make_jwt("power", user_id=USER_ID)
|
||||
req_id = "my-unique-req-id"
|
||||
|
||||
async def _stream(user_id, message, context):
|
||||
async def _stream(user_id, message, context, db_session_factory=None):
|
||||
yield "token", "ok"
|
||||
yield "mutations", []
|
||||
|
||||
with patch("app.api.routes.device_ws.run_home_stream", side_effect=_stream):
|
||||
with client.websocket_connect(f"/api/v1/ws/device?token={token}") as ws:
|
||||
|
||||
Reference in New Issue
Block a user