diff --git a/CLAUDE.md b/CLAUDE.md index 82ec4fc..c9ff607 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -8,16 +8,16 @@ Guidance for working in this repository. ## Layout -All application code lives under `src/`. Packages keep their top-level names (`agents`, `db`, `utils`), so imports are `from agents...` / `from db...` / `from utils...`, with `src` on the path (pytest sets `pythonpath = ["src"]`; the Docker image runs uvicorn with `--app-dir src`). Run entrypoints from the repo root so `load_dotenv()` finds `.env`. +All application code lives under `src/`. Packages keep their top-level names (`agents`, `db`, `emails`), so imports are `from agents...` / `from db...` / `from emails...`, with `src` on the path (pytest sets `pythonpath = ["src"]`; the Docker image runs uvicorn with `--app-dir src`). Run entrypoints from the repo root so `load_dotenv()` finds `.env`. - `src/api.py` — FastAPI service. `POST /run` (full campaign) and `POST /test` (one user). Both run in the background, return `202`, and default to dry-run. Auth via the `X-API-Key` header matched against `SECRET_KEY`. - `src/app.py` — local CLI mirroring the API (`run`, `test`), for testing without HTTP. - `src/agents/orchestrator.py` — the pipeline: analyst → 5 parallel beat desks (researcher → writer → editor) → managing-editor gap roundtable → masthead → reviewer → deterministic clean/assemble/dedupe. Most non-agent logic (citation gating, URL/article validation, dedupe, subject-name canonicalization, prose humanizing) lives here. -- `src/agents/` — one module per agent (`analyst`, `researcher`, `writer`, `editor`, `managing_editor`, `reviewer`), plus `beats.py` (beat desks), `campaign.py` (top-level run over subscriptions), `providers/` (subject-memory and ticker-profile context providers), and `tools/` (Serper search, web fetch). +- `src/agents/` — one module per agent (`analyst`, `researcher`, `writer`, `editor`, `managing_editor`, `reviewer`), plus `beats.py` (beat desks), `campaign.py` (top-level run over subscriptions), `sections.py` (the five editorial beats), `providers/` (subject-memory and ticker-profile context providers), and `tools/` (Serper search, web fetch). +- `src/agents/runtime/` — agent plumbing shared by every agent: `chat_client.py` (per-role chat client plus the `SKILLS` provider), `make_agent.py` (the factory that wires generic activity tracking into every agent), `guardrails.py` (guardrail/citation middleware), and `tracking.py` (the generic `ActivityTracker`/`ToolTracker` middleware, `newsletter_scope`, and run context vars). New agents are built via `make_agent(...)` so tracking is automatic. - `src/agents/skills/` — `SKILL.md` files that control agent behavior (`subject-profile`, `section-research`, `newsletter-format`). Prefer editing these over code when changing how agents research or write. -- `src/db/` — all database access. `mediapulse.py` reads subscriptions and ticker profiles from the upstream MediaPulse Postgres (`MEDIAPULSE_DATABASE_URL`, read-only, raw psycopg). The app's own Postgres (`DATABASE_URL`, SQLModel) uses `engine.py` for the shared engine, with each table's model alongside its operations in `newsletters.py` (archives each generated newsletter as markdown plus JSONB metadata) and `memory.py` (subject-brief agent memory). Tables are auto-created on first write. +- `src/db/` — all database access. `mediapulse.py` reads subscriptions and ticker profiles from the upstream MediaPulse Postgres (`MEDIAPULSE_DATABASE_URL`, read-only, raw psycopg). The app's own Postgres (`DATABASE_URL`, SQLModel) uses `engine.py` for the shared engine, with each table's model alongside its operations: `newsletters.py` (archives each newsletter as markdown plus JSONB metadata, with a `pending`/`complete`/`failed` lifecycle via `create_newsletter`/`finalize_newsletter`), `memory.py` (subject-brief agent memory), and `agent_activity.py` (one row per agent run and tool call, tied to its `newsletter_id`, recording status, duration, model, and token usage). The schema is owned by Alembic migrations, not `create_all` (see Migrations). - `src/emails/` — everything email-related. `mailer.py` sends via Resend. `templates/` pairs each template's renderer with its tokenized HTML: `templates/newsletter.py` parses the newsletter markdown and fills `templates/newsletter.html` (the gitignored build artifact from `email-playground`). -- `src/utils/` — `client.py`, `guardrails.py`, `sections.py`. - `email-playground/` — a standalone React Email (TypeScript) project that is the visual source-of-truth for MediaPulse email templates (the newsletter is the first one). `npm run build:templates` renders each template to a tokenized `src/emails/templates/.html` that the Python side consumes (`emails/templates/newsletter.py` for the newsletter). Those HTML files are gitignored build artifacts, regenerated fresh in CI and the Docker build. Restyle emails there, not in Python. See `email-playground/README.md`. - `tests/` — pytest suite covering the deterministic logic (see Testing below). @@ -39,6 +39,19 @@ python src/app.py test --email=you@example.com # dry-run one user python src/api.py # serve at http://localhost:8000 (docs at /docs) ``` +## Migrations + +The app's own Postgres schema (`DATABASE_URL`) is managed by Alembic, not `create_all`. Migrations live in `alembic/versions/`; `alembic/env.py` puts `src` on the path, loads `.env`, resolves `DATABASE_URL` (reusing `db.engine._engine_url`), and targets `SQLModel.metadata`. + +``` +alembic upgrade head # apply pending migrations +alembic revision -m "add X" # new (hand-written) migration +alembic revision --autogenerate -m "add X" # diff models vs DB, needs a live DATABASE_URL +alembic downgrade -1 # roll back one +``` + +The Docker image runs `alembic upgrade head` on container start via `docker-entrypoint.sh`, then execs the `CMD`. It is skipped when `DATABASE_URL` is unset or `RUN_MIGRATIONS=0`. To run migrations as a standalone release job with the same image: `docker run -e RUN_MIGRATIONS=0 alembic upgrade head`. After changing a SQLModel table, add a migration in the same change so deploys stay in sync. + ## Code quality CI (`.github/workflows/code-quality.yml`) runs exactly these two commands, both must pass: @@ -67,7 +80,7 @@ cd email-playground && npm install && npm run build:templates Config is in `pyproject.toml` under `[tool.pytest.ini_options]`: `pythonpath = ["src"]`, `testpaths = ["tests"]`, `asyncio_mode = "auto"` (async tests need no decorator). CI runs `pytest` in a separate `tests` job after installing both requirement files. -The suite covers the deterministic logic, not the LLM agents: orchestrator text/section helpers, the email template, Serper tools, the MediaPulse profile shaping, the newsletter store and subject-memory upsert, guardrail middleware, the subject-memory and ticker-profile context providers, the mailer, the client model resolution, and the campaign delivery flow. Conventions for new tests: +The suite covers the deterministic logic, not the LLM agents: orchestrator text/section helpers, the email template, Serper tools, the MediaPulse profile shaping, the newsletter store and lifecycle and subject-memory upsert, the agent-activity store, the guardrail and activity-tracking middleware, the subject-memory and ticker-profile context providers, the mailer, the client model resolution, and the campaign delivery flow. Conventions for new tests: - `tests/conftest.py` sets dummy credentials before collection, because importing the agent modules constructs the whole agent graph at import time. - Every external call (httpx, psycopg, the LLM clients) is monkeypatched, and the SQLModel store is exercised against in-memory SQLite. The suite is fully offline. A test that reaches the network is a test bug. diff --git a/Dockerfile b/Dockerfile index e7fc91d..effcc95 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,9 +21,12 @@ COPY . . COPY --from=template-builder /build/src/emails/templates/ src/emails/templates/ RUN useradd --no-create-home --shell /bin/false app \ + && chmod +x docker-entrypoint.sh \ && chown -R app:app /app USER app EXPOSE 8000 +# Entrypoint applies migrations (alembic upgrade head) before launching the command. +ENTRYPOINT ["./docker-entrypoint.sh"] CMD ["uvicorn", "api:app", "--app-dir", "src", "--host", "0.0.0.0", "--port", "8000"] diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 0000000..b8b3c3c --- /dev/null +++ b/alembic.ini @@ -0,0 +1,40 @@ +# Alembic configuration. The database URL is resolved from DATABASE_URL in alembic/env.py, +# so sqlalchemy.url is intentionally left blank here. +[alembic] +script_location = alembic +prepend_sys_path = . +sqlalchemy.url = + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 0000000..bd98442 --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,64 @@ +import os +import sys +from logging.config import fileConfig +from pathlib import Path + +from alembic import context +from sqlalchemy import create_engine + +# Put `src` on the path so the app's db package imports the same way it does at runtime. +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from dotenv import load_dotenv +from sqlmodel import SQLModel + +from db.agent_activity import AgentActivity +from db.engine import _engine_url +from db.memory import SubjectMemory +from db.newsletters import Newsletter + +load_dotenv() + +# Register every model on SQLModel.metadata so autogenerate can see the full schema. +_REGISTERED_MODELS = (Newsletter, SubjectMemory, AgentActivity) +target_metadata = SQLModel.metadata + +config = context.config + +if config.config_file_name is not None: + fileConfig(config.config_file_name) + + +def _database_url() -> str: + if not os.getenv("DATABASE_URL"): + raise RuntimeError("DATABASE_URL is not set; Alembic needs it to run migrations.") + + return _engine_url() + + +def run_migrations_offline() -> None: + context.configure( + url=_database_url(), + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + engine = create_engine(_database_url(), pool_pre_ping=True) + + with engine.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 0000000..8dab86e --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,27 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: str | None = ${repr(down_revision)} +branch_labels: str | Sequence[str] | None = ${repr(branch_labels)} +depends_on: str | Sequence[str] | None = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/0001_initial.py b/alembic/versions/0001_initial.py new file mode 100644 index 0000000..6973698 --- /dev/null +++ b/alembic/versions/0001_initial.py @@ -0,0 +1,74 @@ +"""initial schema: newsletters, subject_memory, agent_activity + +Revision ID: 0001_initial +Revises: +Create Date: 2026-06-18 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "0001_initial" +down_revision: str | None = None +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "newsletters", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("subject", sa.String(), nullable=False), + sa.Column("content", sa.String(), nullable=False), + sa.Column("status", sa.String(), nullable=False), + sa.Column("metadata", postgresql.JSONB(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + ) + op.create_index("ix_newsletters_subject", "newsletters", ["subject"]) + op.create_index("ix_newsletters_status", "newsletters", ["status"]) + + op.create_table( + "subject_memory", + sa.Column("subject", sa.String(), primary_key=True), + sa.Column("brief", sa.String(), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + ) + + op.create_table( + "agent_activity", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("newsletter_id", sa.Integer(), nullable=False), + sa.Column("kind", sa.String(), nullable=False), + sa.Column("name", sa.String(), nullable=False), + sa.Column("model", sa.String(), nullable=True), + sa.Column("status", sa.String(), nullable=False), + sa.Column("duration_ms", sa.Integer(), nullable=True), + sa.Column("input_tokens", sa.Integer(), nullable=True), + sa.Column("output_tokens", sa.Integer(), nullable=True), + sa.Column("total_tokens", sa.Integer(), nullable=True), + sa.Column("error", sa.String(), nullable=True), + sa.Column("metadata", postgresql.JSONB(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + ) + op.create_index("ix_agent_activity_newsletter_id", "agent_activity", ["newsletter_id"]) + op.create_index("ix_agent_activity_name", "agent_activity", ["name"]) + op.create_index("ix_agent_activity_model", "agent_activity", ["model"]) + + +def downgrade() -> None: + op.drop_index("ix_agent_activity_model", table_name="agent_activity") + op.drop_index("ix_agent_activity_name", table_name="agent_activity") + op.drop_index("ix_agent_activity_newsletter_id", table_name="agent_activity") + op.drop_table("agent_activity") + + op.drop_table("subject_memory") + + op.drop_index("ix_newsletters_status", table_name="newsletters") + op.drop_index("ix_newsletters_subject", table_name="newsletters") + op.drop_table("newsletters") diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh new file mode 100755 index 0000000..88f52a7 --- /dev/null +++ b/docker-entrypoint.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env sh +set -e + +# Apply pending migrations on startup, unless disabled (RUN_MIGRATIONS=0) or no database is configured. +# Set RUN_MIGRATIONS=0 and pass `alembic upgrade head` as the command to run migrations as a standalone job. +if [ "${RUN_MIGRATIONS:-1}" != "0" ] && [ -n "$DATABASE_URL" ]; then + echo "Applying database migrations (alembic upgrade head)..." + alembic upgrade head +fi + +exec "$@" diff --git a/environment.yml b/environment.yml index e61c35f..f58d056 100644 --- a/environment.yml +++ b/environment.yml @@ -10,5 +10,6 @@ dependencies: - httpx - psycopg[binary] - sqlmodel + - alembic - fastapi - uvicorn diff --git a/pyproject.toml b/pyproject.toml index fc42c47..393297e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ select = ["E", "F", "I", "W"] [tool.ruff.lint.per-file-ignores] "src/api.py" = ["E402"] "src/app.py" = ["E402"] +"alembic/env.py" = ["E402", "I001"] [tool.pytest.ini_options] pythonpath = ["src"] diff --git a/requirements.txt b/requirements.txt index 7850382..b8d44a2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,5 +3,6 @@ python-dotenv httpx psycopg[binary] sqlmodel +alembic fastapi uvicorn diff --git a/src/agents/analyst.py b/src/agents/analyst.py index e069958..ccdcb42 100644 --- a/src/agents/analyst.py +++ b/src/agents/analyst.py @@ -1,12 +1,11 @@ -from agent_framework import Agent - from agents.providers.memory import SubjectMemoryProvider from agents.providers.ticker import TickerProfileProvider +from agents.runtime.chat_client import SKILLS, chat_client +from agents.runtime.guardrails import SubjectGuardrail +from agents.runtime.make_agent import make_agent from agents.tools import search, web_fetch -from utils.client import SKILLS, chat_client -from utils.guardrails import SubjectGuardrail -analyst = Agent( +analyst = make_agent( name="analyst", description="Resolves a subject into a research brief.", client=chat_client("analyst"), diff --git a/src/agents/beats.py b/src/agents/beats.py index dc3cc13..a4100f4 100644 --- a/src/agents/beats.py +++ b/src/agents/beats.py @@ -3,9 +3,9 @@ from agent_framework import Agent from agents.researcher import make_researcher +from agents.runtime.guardrails import SourceRegistry +from agents.sections import SECTIONS, Section from agents.writer import make_writer -from utils.guardrails import SourceRegistry -from utils.sections import SECTIONS, Section @dataclass diff --git a/src/agents/campaign.py b/src/agents/campaign.py index c86ec0c..59c0cc1 100644 --- a/src/agents/campaign.py +++ b/src/agents/campaign.py @@ -2,8 +2,9 @@ import time from agents.orchestrator import run_newsletter +from agents.runtime.tracking import newsletter_scope from db.mediapulse import fetch_subscriptions -from db.newsletters import save_newsletter +from db.newsletters import create_newsletter, finalize_newsletter from emails.mailer import send_email from emails.templates.newsletter import newsletter_sources, render_newsletter_email @@ -29,10 +30,12 @@ async def run_campaign(*, subscriptions: list[dict] | None = None, send: bool = last_send = 0.0 delivered: list[dict] = [] - async def deliver(ticker: str, recipients: list[dict], markdown: str) -> None: + async def deliver(ticker: str, recipients: list[dict], markdown: str, newsletter_id: int | None) -> None: nonlocal last_send email = render_newsletter_email(markdown, ticker=ticker) - save_newsletter(ticker, markdown, {"ticker": ticker, "sources": newsletter_sources(markdown)}) + finalize_newsletter( + newsletter_id, content=markdown, metadata={"ticker": ticker, "sources": newsletter_sources(markdown)} + ) for subscription in recipients: if send: @@ -53,19 +56,23 @@ async def deliver(ticker: str, recipients: list[dict], markdown: str) -> None: delivered.append({"email": subscription["email"], "ticker": ticker}) async def process(ticker: str, recipients: list[dict]) -> None: - async with generate_semaphore: - log(f"generating {ticker} ...") + newsletter_id = create_newsletter(ticker) - try: - markdown = await run_newsletter(ticker) - except Exception as error: - log(f"failed {ticker}: {error}") + with newsletter_scope(newsletter_id): + async with generate_semaphore: + log(f"generating {ticker} ...") - return + try: + markdown = await run_newsletter(ticker) + except Exception as error: + finalize_newsletter(newsletter_id, status="failed") + log(f"failed {ticker}: {error}") - log(f"done {ticker}") + return - await deliver(ticker, recipients, markdown) + log(f"done {ticker}") + + await deliver(ticker, recipients, markdown, newsletter_id) await asyncio.gather(*(process(ticker, recipients) for ticker, recipients in subscribers_by_ticker.items())) diff --git a/src/agents/editor.py b/src/agents/editor.py index 2c03a5d..7e2267d 100644 --- a/src/agents/editor.py +++ b/src/agents/editor.py @@ -1,8 +1,7 @@ -from agent_framework import Agent +from agents.runtime.chat_client import SKILLS, chat_client +from agents.runtime.make_agent import make_agent -from utils.client import SKILLS, chat_client - -editor = Agent( +editor = make_agent( name="editor", description="Writes the masthead and reviews sections.", client=chat_client("editor"), diff --git a/src/agents/managing_editor.py b/src/agents/managing_editor.py index 60a38a9..f0d3df9 100644 --- a/src/agents/managing_editor.py +++ b/src/agents/managing_editor.py @@ -1,8 +1,7 @@ -from agent_framework import Agent +from agents.runtime.chat_client import SKILLS, chat_client +from agents.runtime.make_agent import make_agent -from utils.client import SKILLS, chat_client - -managing_editor = Agent( +managing_editor = make_agent( name="managing_editor", description="Chairs the newsroom roundtable and finds coverage gaps across the edition.", client=chat_client("managing_editor"), diff --git a/src/agents/orchestrator.py b/src/agents/orchestrator.py index fe9a0f5..9acb9b9 100644 --- a/src/agents/orchestrator.py +++ b/src/agents/orchestrator.py @@ -7,8 +7,8 @@ from agents.editor import editor from agents.managing_editor import managing_editor from agents.reviewer import reviewer +from agents.sections import SECTIONS from db.memory import remember_subject -from utils.sections import SECTIONS _URL_RE = re.compile(r"https?://[^\s\)]+") _READ_RE = re.compile(r"\[Read:[^\]]*\]\(([^)]*)\)") diff --git a/src/agents/researcher.py b/src/agents/researcher.py index 6abb041..ec72d10 100644 --- a/src/agents/researcher.py +++ b/src/agents/researcher.py @@ -1,14 +1,15 @@ from agent_framework import Agent +from agents.runtime.chat_client import SKILLS, chat_client +from agents.runtime.guardrails import RecordSources, SourceRegistry +from agents.runtime.make_agent import make_agent +from agents.sections import Section from agents.tools import search -from utils.client import SKILLS, chat_client -from utils.guardrails import RecordSources, SourceRegistry -from utils.sections import Section def make_researcher(section: Section, registry: SourceRegistry) -> Agent: """A news researcher for one beat: finds and ranks candidate articles.""" - return Agent( + return make_agent( name=f"researcher_{section.slug}", description=f"Finds and ranks sources for the {section.name} beat.", client=chat_client("researcher"), diff --git a/src/agents/reviewer.py b/src/agents/reviewer.py index 8e46c3b..3a8719d 100644 --- a/src/agents/reviewer.py +++ b/src/agents/reviewer.py @@ -1,8 +1,7 @@ -from agent_framework import Agent +from agents.runtime.chat_client import chat_client +from agents.runtime.make_agent import make_agent -from utils.client import chat_client - -reviewer = Agent( +reviewer = make_agent( name="reviewer", description="Critiques the whole edition for quality before publication.", client=chat_client("reviewer"), diff --git a/src/utils/__init__.py b/src/agents/runtime/__init__.py similarity index 100% rename from src/utils/__init__.py rename to src/agents/runtime/__init__.py diff --git a/src/utils/client.py b/src/agents/runtime/chat_client.py similarity index 64% rename from src/utils/client.py rename to src/agents/runtime/chat_client.py index 0343034..a0e9050 100644 --- a/src/utils/client.py +++ b/src/agents/runtime/chat_client.py @@ -4,15 +4,11 @@ from agent_framework import SkillsProvider from agent_framework.openai import OpenAIChatCompletionClient -SKILLS = SkillsProvider.from_paths(skill_paths=Path(__file__).parent.parent / "agents" / "skills") +SKILLS = SkillsProvider.from_paths(skill_paths=Path(__file__).parent.parent / "skills") def chat_client(role: str) -> OpenAIChatCompletionClient: - """Build a chat client for an agent role. - - The model is resolved from `_MODEL`, then `OPENAI_MODEL`, then a default, - so each agent can run its own model while sharing one base URL and key. - """ + """Build a chat client for an agent role (model from `_MODEL`, then `OPENAI_MODEL`, then a default).""" model = os.getenv(f"{role.upper()}_MODEL") or os.getenv("OPENAI_MODEL", "gpt-4.1-mini") return OpenAIChatCompletionClient(model=model, base_url=os.getenv("OPENAI_BASE_URL")) diff --git a/src/utils/guardrails.py b/src/agents/runtime/guardrails.py similarity index 100% rename from src/utils/guardrails.py rename to src/agents/runtime/guardrails.py diff --git a/src/agents/runtime/make_agent.py b/src/agents/runtime/make_agent.py new file mode 100644 index 0000000..0dd2e74 --- /dev/null +++ b/src/agents/runtime/make_agent.py @@ -0,0 +1,18 @@ +from typing import Any + +from agent_framework import Agent + +from agents.runtime.tracking import ACTIVITY_TRACKER, TOOL_TRACKER + + +def make_agent(**kwargs: Any) -> Agent: + """Build an Agent with tracking wired in: a shared ActivityTracker, plus a ToolTracker when it has tools.""" + extra_middleware = list(kwargs.pop("middleware", None) or []) + trackers: list[Any] = [ACTIVITY_TRACKER] + + if kwargs.get("tools"): + trackers.append(TOOL_TRACKER) + + kwargs["middleware"] = trackers + extra_middleware + + return Agent(**kwargs) diff --git a/src/agents/runtime/tracking.py b/src/agents/runtime/tracking.py new file mode 100644 index 0000000..389d9ac --- /dev/null +++ b/src/agents/runtime/tracking.py @@ -0,0 +1,155 @@ +import asyncio +import contextvars +import time +from collections.abc import Awaitable, Callable, Mapping +from contextlib import contextmanager + +from agent_framework import ( + AgentContext, + AgentMiddleware, + FunctionInvocationContext, + FunctionMiddleware, + MiddlewareTermination, +) + +from db.agent_activity import log_activity + +# The id of the newsletter every agent/tool event in the current run belongs to (None outside a tracked run). +current_newsletter_id: contextvars.ContextVar[int | None] = contextvars.ContextVar( + "current_newsletter_id", default=None +) +# The name of the agent currently executing, so tool events can be attributed to their agent. +current_agent_name: contextvars.ContextVar[str | None] = contextvars.ContextVar("current_agent_name", default=None) + + +@contextmanager +def newsletter_scope(newsletter_id: int | None): + """Bind every agent/tool event inside the block to `newsletter_id` (propagates into asyncio tasks).""" + token = current_newsletter_id.set(newsletter_id) + + try: + yield newsletter_id + finally: + current_newsletter_id.reset(token) + + +def _elapsed_ms(start_time: float) -> int: + return int((time.monotonic() - start_time) * 1000) + + +def _serialize_arguments(arguments: object) -> object: + """Reduce tool-call arguments to a JSON-serializable form for the activity `meta` column.""" + if arguments is None: + return None + + if hasattr(arguments, "model_dump"): + try: + return arguments.model_dump(mode="json") + except Exception: + return str(arguments) + + if isinstance(arguments, Mapping): + return dict(arguments) + + return str(arguments) + + +def _result_text(result: object) -> str: + if result is None: + return "" + + if isinstance(result, str): + return result + + if isinstance(result, (list, tuple)): + return " ".join(_result_text(item) for item in result) + + return getattr(result, "text", None) or str(result) + + +async def _log_event(newsletter_id: int | None, **event_fields: object) -> None: + """Persist one activity event off the event loop. Best-effort: tracking never breaks a run.""" + if newsletter_id is None: + return + + try: + await asyncio.to_thread(log_activity, newsletter_id=newsletter_id, **event_fields) + except Exception: + pass + + +class ActivityTracker(AgentMiddleware): + """Record every agent run: status, duration, and token usage, tied to the current newsletter.""" + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + newsletter_id = current_newsletter_id.get() + agent_name = getattr(context.agent, "name", None) or getattr(context.agent, "id", None) or "agent" + client = getattr(context.agent, "client", None) + model = getattr(client, "model", None) + agent_name_token = current_agent_name.set(agent_name) + start_time = time.monotonic() + status, error = "ok", None + + try: + await call_next() + except MiddlewareTermination: + status = "terminated" + + raise + except Exception as exception: + status, error = "error", f"{type(exception).__name__}: {exception}" + + raise + finally: + current_agent_name.reset(agent_name_token) + result = context.result + usage_details = getattr(result, "usage_details", None) or {} + + await _log_event( + newsletter_id, + kind="agent", + name=agent_name, + model=model, + status=status, + duration_ms=_elapsed_ms(start_time), + input_tokens=usage_details.get("input_token_count"), + output_tokens=usage_details.get("output_token_count"), + total_tokens=usage_details.get("total_token_count"), + error=error, + meta={"finish_reason": getattr(result, "finish_reason", None)}, + ) + + +class ToolTracker(FunctionMiddleware): + """Record every tool call: status, duration, arguments, and the owning agent.""" + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + newsletter_id = current_newsletter_id.get() + tool_name = getattr(context.function, "name", None) or "tool" + start_time = time.monotonic() + status, error = "ok", None + + try: + await call_next() + except Exception as exception: + status, error = "error", f"{type(exception).__name__}: {exception}" + + raise + finally: + await _log_event( + newsletter_id, + kind="tool", + name=tool_name, + status=status, + duration_ms=_elapsed_ms(start_time), + error=error, + meta={ + "agent": current_agent_name.get(), + "arguments": _serialize_arguments(context.arguments), + "result_chars": len(_result_text(context.result)), + }, + ) + + +ACTIVITY_TRACKER = ActivityTracker() +TOOL_TRACKER = ToolTracker() diff --git a/src/utils/sections.py b/src/agents/sections.py similarity index 100% rename from src/utils/sections.py rename to src/agents/sections.py diff --git a/src/agents/writer.py b/src/agents/writer.py index de5cb36..0a3908f 100644 --- a/src/agents/writer.py +++ b/src/agents/writer.py @@ -1,14 +1,15 @@ from agent_framework import Agent +from agents.runtime.chat_client import SKILLS, chat_client +from agents.runtime.guardrails import EnforceCitations, SourceRegistry +from agents.runtime.make_agent import make_agent +from agents.sections import Section from agents.tools import web_fetch -from utils.client import SKILLS, chat_client -from utils.guardrails import EnforceCitations, SourceRegistry -from utils.sections import Section def make_writer(section: Section, registry: SourceRegistry) -> Agent: """Writes one beat's section from the researcher's candidate articles.""" - return Agent( + return make_agent( name=f"writer_{section.slug}", description=f"Writes the {section.name} section from researched sources.", client=chat_client("writer"), diff --git a/src/db/agent_activity.py b/src/db/agent_activity.py new file mode 100644 index 0000000..451f4c9 --- /dev/null +++ b/src/db/agent_activity.py @@ -0,0 +1,78 @@ +import logging +from datetime import datetime + +from sqlalchemy import JSON, Column, DateTime, func +from sqlalchemy.dialects.postgresql import JSONB +from sqlmodel import Field, Session, SQLModel + +from db.engine import get_engine, is_configured + +logger = logging.getLogger(__name__) + +# JSONB on Postgres, plain JSON elsewhere so the offline test suite can use SQLite. +_JSON = JSON().with_variant(JSONB(), "postgresql") + + +class AgentActivity(SQLModel, table=True): + __tablename__ = "agent_activity" + + id: int | None = Field(default=None, primary_key=True) + newsletter_id: int = Field(index=True) + kind: str # "agent" | "tool" + name: str = Field(index=True) + model: str | None = Field(default=None, index=True) # the model behind an agent run (null for tools) + status: str # "ok" | "error" | "terminated" + duration_ms: int | None = None + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + error: str | None = None + meta: dict = Field(default_factory=dict, sa_column=Column("metadata", _JSON, nullable=False)) + created_at: datetime | None = Field( + default=None, sa_column=Column(DateTime(timezone=True), nullable=False, server_default=func.now()) + ) + + +def log_activity( + *, + newsletter_id: int, + kind: str, + name: str, + status: str, + model: str | None = None, + duration_ms: int | None = None, + input_tokens: int | None = None, + output_tokens: int | None = None, + total_tokens: int | None = None, + error: str | None = None, + meta: dict | None = None, +) -> int | None: + """Record one agent or tool event and return its row id, or None if storage is unconfigured or the write fails.""" + if not is_configured(): + return None + + try: + activity = AgentActivity( + newsletter_id=newsletter_id, + kind=kind, + name=name, + status=status, + model=model, + duration_ms=duration_ms, + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + error=error, + meta=meta or {}, + ) + + with Session(get_engine()) as session: + session.add(activity) + session.commit() + session.refresh(activity) + + return activity.id + except Exception as write_error: + logger.warning("failed to log activity %s/%s: %s", kind, name, write_error) + + return None diff --git a/src/db/engine.py b/src/db/engine.py index 1c84ec4..c5ec3da 100644 --- a/src/db/engine.py +++ b/src/db/engine.py @@ -1,6 +1,6 @@ import os -from sqlmodel import SQLModel, create_engine +from sqlmodel import create_engine _engine = None @@ -22,16 +22,10 @@ def _engine_url() -> str: def get_engine(): - """Build the engine once and create the app's tables (idempotent, once per process).""" + """Build the engine once per process. The schema is owned by Alembic migrations, not create_all.""" global _engine if _engine is None: - # Import here so the models register on SQLModel.metadata before create_all, avoiding a cycle. - from db.memory import SubjectMemory - from db.newsletters import Newsletter - - engine = create_engine(_engine_url(), pool_pre_ping=True) - SQLModel.metadata.create_all(engine, tables=[Newsletter.__table__, SubjectMemory.__table__]) - _engine = engine + _engine = create_engine(_engine_url(), pool_pre_ping=True) return _engine diff --git a/src/db/newsletters.py b/src/db/newsletters.py index 621cd79..a00a80b 100644 --- a/src/db/newsletters.py +++ b/src/db/newsletters.py @@ -18,20 +18,21 @@ class Newsletter(SQLModel, table=True): id: int | None = Field(default=None, primary_key=True) subject: str = Field(index=True) - content: str + content: str = "" + status: str = Field(default="complete", index=True) # "pending" | "complete" | "failed" meta: dict = Field(default_factory=dict, sa_column=Column("metadata", _JSON, nullable=False)) created_at: datetime | None = Field( default=None, sa_column=Column(DateTime(timezone=True), nullable=False, server_default=func.now()) ) -def save_newsletter(subject: str, content: str, metadata: dict | None = None) -> int | None: - """Archive a generated newsletter and return its row id, or None if storage is unconfigured or the write fails.""" +def create_newsletter(subject: str) -> int | None: + """Open a placeholder newsletter row before generation and return its id, or None if storage is unconfigured.""" if not is_configured(): return None try: - newsletter = Newsletter(subject=subject, content=content, meta=metadata or {}) + newsletter = Newsletter(subject=subject, content="", status="pending") with Session(get_engine()) as session: session.add(newsletter) @@ -40,6 +41,45 @@ def save_newsletter(subject: str, content: str, metadata: dict | None = None) -> return newsletter.id except Exception as error: - logger.warning("failed to store newsletter for %s: %s", subject, error) + logger.warning("failed to open newsletter for %s: %s", subject, error) + + return None + + +def finalize_newsletter( + newsletter_id: int | None, *, content: str | None = None, metadata: dict | None = None, status: str = "complete" +) -> None: + """Fill in a previously opened newsletter row (best-effort, no-op if unconfigured or the row is missing).""" + if newsletter_id is None or not is_configured(): + return + + try: + with Session(get_engine()) as session: + newsletter = session.get(Newsletter, newsletter_id) + + if newsletter is None: + return + + if content is not None: + newsletter.content = content + + if metadata is not None: + newsletter.meta = metadata + + newsletter.status = status + session.add(newsletter) + session.commit() + except Exception as error: + logger.warning("failed to finalize newsletter %s: %s", newsletter_id, error) + +def save_newsletter(subject: str, content: str, metadata: dict | None = None) -> int | None: + """Archive a generated newsletter and return its row id, or None if storage is unconfigured or the write fails.""" + newsletter_id = create_newsletter(subject) + + if newsletter_id is None: return None + + finalize_newsletter(newsletter_id, content=content, metadata=metadata or {}) + + return newsletter_id diff --git a/tests/test_agent_activity.py b/tests/test_agent_activity.py new file mode 100644 index 0000000..1fd9191 --- /dev/null +++ b/tests/test_agent_activity.py @@ -0,0 +1,72 @@ +from sqlalchemy.pool import StaticPool +from sqlmodel import Session, SQLModel, create_engine, select + +import db.agent_activity as agent_activity +import db.engine as engine_module + + +def _use_sqlite(monkeypatch): + engine = create_engine("sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool) + SQLModel.metadata.create_all(engine, tables=[agent_activity.AgentActivity.__table__]) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@h/db") + monkeypatch.setattr(engine_module, "_engine", engine) + + return engine + + +def test_log_activity_skips_without_env(monkeypatch): + monkeypatch.delenv("DATABASE_URL", raising=False) + + assert agent_activity.log_activity(newsletter_id=1, kind="agent", name="analyst", status="ok") is None + + +def test_log_activity_inserts_and_roundtrips(monkeypatch): + engine = _use_sqlite(monkeypatch) + row_id = agent_activity.log_activity( + newsletter_id=7, + kind="agent", + name="researcher_quick_hits", + model="gpt-4.1-mini", + status="ok", + duration_ms=1200, + input_tokens=300, + output_tokens=120, + total_tokens=420, + meta={"finish_reason": "stop"}, + ) + + assert isinstance(row_id, int) + + with Session(engine) as session: + rows = session.exec(select(agent_activity.AgentActivity)).all() + + assert len(rows) == 1 + assert rows[0].newsletter_id == 7 + assert rows[0].kind == "agent" + assert rows[0].name == "researcher_quick_hits" + assert rows[0].model == "gpt-4.1-mini" + assert rows[0].total_tokens == 420 + assert rows[0].meta == {"finish_reason": "stop"} + + +def test_log_activity_records_failure(monkeypatch): + engine = _use_sqlite(monkeypatch) + agent_activity.log_activity(newsletter_id=9, kind="tool", name="search", status="error", error="RuntimeError: boom") + + with Session(engine) as session: + row = session.exec(select(agent_activity.AgentActivity)).one() + + assert row.status == "error" + assert row.error == "RuntimeError: boom" + assert row.total_tokens is None + + +def test_log_activity_swallows_errors(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@h/db") + + def boom(): + raise RuntimeError("db down") + + monkeypatch.setattr(agent_activity, "get_engine", boom) + + assert agent_activity.log_activity(newsletter_id=1, kind="agent", name="analyst", status="ok") is None diff --git a/tests/test_campaign.py b/tests/test_campaign.py index b50f8e8..ee449e5 100644 --- a/tests/test_campaign.py +++ b/tests/test_campaign.py @@ -67,19 +67,37 @@ async def test_run_campaign_fetches_subscriptions_when_unset(monkeypatch): assert result["delivered"] == [{"email": "z@x.com", "ticker": "ACME"}] -async def test_each_ticker_is_saved_once_with_ticker_and_sources(monkeypatch): +async def test_each_ticker_is_finalized_once_with_ticker_and_sources(monkeypatch): _patch_pipeline(monkeypatch) - saved = [] + finalized = [] - def fake_save(subject, markdown, metadata): - saved.append((subject, metadata)) + def fake_finalize(newsletter_id, *, content=None, metadata=None, status="complete"): + finalized.append((metadata, status)) - monkeypatch.setattr(campaign, "save_newsletter", fake_save) + monkeypatch.setattr(campaign, "create_newsletter", lambda subject: 1) + monkeypatch.setattr(campaign, "finalize_newsletter", fake_finalize) await campaign.run_campaign(subscriptions=_subs(), send=False, log=lambda *a: None) - assert sorted(subject for subject, _meta in saved) == ["ACME", "GLOBEX"] # once per unique ticker + completed = [metadata for metadata, _status in finalized if metadata is not None] - for _subject, metadata in saved: - assert metadata["ticker"] in {"ACME", "GLOBEX"} + assert sorted(metadata["ticker"] for metadata in completed) == ["ACME", "GLOBEX"] # once per unique ticker + + for metadata in completed: assert isinstance(metadata["sources"], list) + + +async def test_failed_ticker_is_finalized_as_failed(monkeypatch): + _patch_pipeline(monkeypatch, fail_for={"ACME"}) + finalized = [] + + def fake_finalize(newsletter_id, *, content=None, metadata=None, status="complete"): + finalized.append(status) + + monkeypatch.setattr(campaign, "create_newsletter", lambda subject: 1) + monkeypatch.setattr(campaign, "finalize_newsletter", fake_finalize) + + await campaign.run_campaign(subscriptions=_subs(), send=False, log=lambda *a: None) + + assert "failed" in finalized # the ACME run was marked failed + assert "complete" in finalized # GLOBEX still completed diff --git a/tests/test_client.py b/tests/test_chat_client.py similarity index 96% rename from tests/test_client.py rename to tests/test_chat_client.py index 80e52cb..6f765b1 100644 --- a/tests/test_client.py +++ b/tests/test_chat_client.py @@ -1,4 +1,4 @@ -import utils.client as client +import agents.runtime.chat_client as client class FakeChatClient: diff --git a/tests/test_guardrails.py b/tests/test_guardrails.py index 4325f1c..703f949 100644 --- a/tests/test_guardrails.py +++ b/tests/test_guardrails.py @@ -1,7 +1,7 @@ import pytest from agent_framework import AgentResponse, Message, MiddlewareTermination -import utils.guardrails as guardrails +import agents.runtime.guardrails as guardrails class Msg: diff --git a/tests/test_newsletters.py b/tests/test_newsletters.py index e14ee4b..7ddb96a 100644 --- a/tests/test_newsletters.py +++ b/tests/test_newsletters.py @@ -45,3 +45,48 @@ def boom(): monkeypatch.setattr(newsletters, "get_engine", boom) assert newsletters.save_newsletter("ACME", "content", {}) is None + + +def test_create_newsletter_opens_pending_row(monkeypatch): + engine = _use_sqlite(monkeypatch) + newsletter_id = newsletters.create_newsletter("ACME") + + assert isinstance(newsletter_id, int) + + with Session(engine) as session: + row = session.get(newsletters.Newsletter, newsletter_id) + + assert row.status == "pending" + assert row.content == "" + + +def test_finalize_newsletter_fills_and_completes(monkeypatch): + engine = _use_sqlite(monkeypatch) + newsletter_id = newsletters.create_newsletter("ACME") + metadata = {"ticker": "ACME", "sources": ["https://x.com/a"]} + newsletters.finalize_newsletter(newsletter_id, content="# ACME Pulse", metadata=metadata) + + with Session(engine) as session: + row = session.get(newsletters.Newsletter, newsletter_id) + + assert row.status == "complete" + assert row.content == "# ACME Pulse" + assert row.meta == metadata + + +def test_finalize_newsletter_marks_failed(monkeypatch): + engine = _use_sqlite(monkeypatch) + newsletter_id = newsletters.create_newsletter("ACME") + newsletters.finalize_newsletter(newsletter_id, status="failed") + + with Session(engine) as session: + row = session.get(newsletters.Newsletter, newsletter_id) + + assert row.status == "failed" + assert row.content == "" + + +def test_finalize_newsletter_noop_without_id(monkeypatch): + _use_sqlite(monkeypatch) + + newsletters.finalize_newsletter(None, content="ignored") # must not raise diff --git a/tests/test_sections.py b/tests/test_sections.py index 169f3f9..066f49d 100644 --- a/tests/test_sections.py +++ b/tests/test_sections.py @@ -1,4 +1,4 @@ -from utils.sections import SECTIONS, Section +from agents.sections import SECTIONS, Section def test_slug_replaces_ampersand_and_spaces(): diff --git a/tests/test_tracking.py b/tests/test_tracking.py new file mode 100644 index 0000000..0280fc7 --- /dev/null +++ b/tests/test_tracking.py @@ -0,0 +1,129 @@ +import pytest +from agent_framework import MiddlewareTermination + +import agents.runtime.tracking as tracking + + +class FakeClient: + model = "gpt-test" + + +class FakeAgent: + name = "analyst" + client = FakeClient() + + +class FakeResult: + def __init__(self, usage=None, finish_reason="stop"): + self.usage_details = usage or {} + self.finish_reason = finish_reason + + +class FakeAgentContext: + def __init__(self): + self.agent = FakeAgent() + self.result = None + + +class FakeFunction: + name = "search" + + +class FakeToolContext: + def __init__(self, arguments, result): + self.function = FakeFunction() + self.arguments = arguments + self.result = result + + +def _capture(monkeypatch): + events = [] + + def fake_log_activity(**kwargs): + events.append(kwargs) + + monkeypatch.setattr(tracking, "log_activity", fake_log_activity) + + return events + + +async def test_activity_tracker_records_ok(monkeypatch): + events = _capture(monkeypatch) + context = FakeAgentContext() + + async def call_next(): + context.result = FakeResult(usage={"input_token_count": 10, "output_token_count": 4, "total_token_count": 14}) + + with tracking.newsletter_scope(42): + await tracking.ActivityTracker().process(context, call_next) + + assert len(events) == 1 + event = events[0] + assert event["newsletter_id"] == 42 + assert event["kind"] == "agent" + assert event["name"] == "analyst" + assert event["model"] == "gpt-test" + assert event["status"] == "ok" + assert event["total_tokens"] == 14 + assert event["meta"] == {"finish_reason": "stop"} + + +async def test_activity_tracker_records_termination(monkeypatch): + events = _capture(monkeypatch) + context = FakeAgentContext() + + async def call_next(): + raise MiddlewareTermination(result=None) + + with tracking.newsletter_scope(1): + with pytest.raises(MiddlewareTermination): + await tracking.ActivityTracker().process(context, call_next) + + assert events[0]["status"] == "terminated" + + +async def test_activity_tracker_records_error(monkeypatch): + events = _capture(monkeypatch) + context = FakeAgentContext() + + async def call_next(): + raise RuntimeError("boom") + + with tracking.newsletter_scope(1): + with pytest.raises(RuntimeError): + await tracking.ActivityTracker().process(context, call_next) + + assert events[0]["status"] == "error" + assert "RuntimeError: boom" == events[0]["error"] + + +async def test_activity_tracker_skips_outside_scope(monkeypatch): + events = _capture(monkeypatch) + context = FakeAgentContext() + + async def call_next(): + context.result = FakeResult() + + await tracking.ActivityTracker().process(context, call_next) # no newsletter_scope active + + assert events == [] + + +async def test_tool_tracker_records_call(monkeypatch): + events = _capture(monkeypatch) + context = FakeToolContext(arguments={"query": "ACME news"}, result="some result text") + + async def call_next(): + return None + + with tracking.newsletter_scope(5): + await tracking.ToolTracker().process(context, call_next) + + assert len(events) == 1 + event = events[0] + assert event["newsletter_id"] == 5 + assert event["kind"] == "tool" + assert event["name"] == "search" + assert event["status"] == "ok" + assert event["meta"]["arguments"] == {"query": "ACME news"} + assert event["meta"]["result_chars"] == len("some result text")