## Implementation Plan: Display Deep Research Costs for Admin Users ### Scope Modify six source files plus tests to extract OpenAI usage headers in the deep search LLM client, propagate cost through the orchestration layer, store it in the worker report, expose it via the router (conditionally for admin users), and render it in the session template. --- ### 1. Extract cost headers in LLM client (`devplacepy/services/deepsearch/llm.py`) **Current (lines 37–44):** ```python response = await client.post(...) ... data = response.json() return data, data.get("usage") or {}, elapsed_ms ``` **Change:** - After `response = await client.post(...)`, capture `response.headers`. - Call `parse_usage_headers(response_headers)` from `devplacepy.services.openai_gateway.usage`. - Return a dict like `{"prompt_tokens": ..., "completion_tokens": ..., "cost_usd": ...}` alongside the existing tuple. - Update the function signature and all callers (in `orchestrate.py` and elsewhere) to accept the additional cost dict. **Signature change example:** ```python async def call_llm(...) -> tuple[dict, dict, float, dict | None]: ... usage_headers = parse_usage_headers(response.headers) return data, body_usage, elapsed_ms, usage_headers ``` --- ### 2. Extract cost headers in query planner (`devplacepy/services/jobs/deepsearch/enhance.py`) **Current (lines 86–92):** ```python response = await client.post(INTERNAL_GATEWAY_URL, ...) data = response.json() ... return data ``` **Change:** - Same pattern as llm.py. After `response.json()`, call `parse_usage_headers(response.headers)`. - Return the cost dict as an additional element. - Update the caller in `orchestrate.py` or wherever the planner result is consumed. --- ### 3. Thread cost through orchestration (`devplacepy/services/jobs/deepsearch/orchestrate.py`) **Current (`_complete` method, lines ~164–168):** Only returns `(text, metrics_dict)`. **Change:** - After calling `call_llm`, merge the usage-header cost from step 1 into the orchestration state. - After calling the planner (step 2), merge that cost as well. - At the end of `_agent_done()`, accumulate total cost across all agents. - Modify the orchestration result to include a `total_cost_usd` field (or a dict of per-agent costs). - Ensure the final report dict (produced in worker.py) receives this cost data. --- ### 4. Add cost fields to worker report (`devplacepy/services/jobs/deepsearch/worker.py`) **Current (lines 215–231, report dict):** `report = { "query": ..., "depth": ..., "score": ..., ..., "chunk_count": ... }` **Change:** - Add `"total_cost_usd": ` to the dict. - If per-agent breakdown is desired, add `"agent_costs": [...]`. - Ensure the value is populated from the orchestration result. --- ### 5. Plumb cost in router (`devplacepy/routers/tools/deepsearch.py`) **Current `_session_context` (lines 225–257):** Returns `viewer_is_admin` but no cost info. **Change:** - Load the report from the DB (already happening in the route). - Extract `total_cost_usd` (and optionally per-agent costs) from the report. - Add to the context dict, but **only if `viewer_is_admin` is True**. - Example: ```python if viewer_is_admin: context["cost_usd"] = report.get("total_cost_usd", 0.0) ``` --- ### 6. Render cost in template (`devplacepy/templates/tools/deepsearch_session.html`) **Current (lines 13–18):** ```html
...
``` **Change:** - Add a conditional block inside the metrics div: ```html {% if viewer_is_admin and cost_usd is defined %} {{ "%.6f"|format(cost_usd) }} cost (USD) {% endif %} ``` --- ### 7. Update tests (`tests/api/tools/deepsearch/session.py`) **Current `_seed_done_session()`:** Seeds report dict without cost fields. **Change:** - Add `"total_cost_usd": 0.012345` to the seed data. - Add an assertion in the session test that for an admin user the rendered page contains the cost string. - For a non-admin user, assert the cost string is absent. --- ### 8. Ensure existing callers of `call_llm` are updated - In `orchestrate.py` and any other file that calls `call_llm` in `devplacepy/services/deepsearch/llm.py`, adjust for the new return signature. - The planner function in `enhance.py` will also have callers that need updating. --- ### Verification Approach Because the system is not fully runnable in this environment, verification will rely on: - **Syntax check** on all modified files: `python3 -m py_compile ` - **Existing unit tests** for the usage header parser: `python3 -m pytest tests/unit/services/openai_gateway/usage.py -x -q` - **Existing unit tests** for deep search (if any): `python3 -m pytest tests/unit/services/deepsearch/ -x -q` - **Lint check** using `ruff check .` (after installing ruff). - **Full unit test suite** (skipping integration tests) using `make test-unit` (requires proper environment). --- ## Definition of Done - [ ] All six source files and the test file have been modified as described above. - [ ] Each modified file passes `python3 -m py_compile` with no errors. - [ ] `python3 -m pytest tests/unit/services/openai_gateway/usage.py -x -q` passes. - [ ] `ruff check .` returns exit code 0 (no lint violations in changed files). - [ ] `make test-unit` (or equivalent) passes all unit tests. - [ ] Adding cost to the report dict does not break any existing endpoint (no regressions). - [ ] The template renders the cost metric for admin users and hides it for non-admins (verified by test assertion if possible, or by manual review of the template logic).