agentflow.py

import traceback
import tomllib
import json
import os

from metaflow import (
    FlowSpec,
    step,
    Config,
    Parameter,
    current,
    retry,
    card,
    secrets,
    pypi,
    pypi_base,
)
from session_service import with_session_service

MAX_AGENT_FAILURES = 3


def parse_csv(entries):
    def _parse():
        for entry in entries.splitlines():
            name, born, synopsis = entry.split(",", 2)
            yield {"ARTIST": name, "BORN": born, "SYNOPSIS": synopsis}

    return {"artists": list(_parse())}


def read_html(html):
    return {"html": html}


@pypi_base(python="3.12")
@with_session_service(name="musicbrainz")
class ResearchAgentFlow(FlowSpec):

    timeline_html = Config(
        "timeline_html", default="timelinecard.html", parser=read_html
    )
    input_list = Config("inputs", default="artists.csv", parser=parse_csv)
    prompts = Config("prompts", default="prompts.toml", parser=tomllib.loads)
    max_artists = Parameter("max-artists", type=int, default=-1)
    checkpoint_interval = Parameter("checkpoint-interval", default=120)

    @step
    def start(self):
        maxx = None if self.max_artists == -1 else self.max_artists
        self.artists = list(map(dict, self.input_list.artists[:maxx]))
        self.next(self.research, foreach="artists")

    @pypi(packages={"openai-agents": "0.2.9", "pydantic": "2.11.7"})
    @secrets(sources=["outerbounds.music-research-agent"])
    @retry
    @step
    def research(self):
        from research_agent import ResearchAgent

        self.artist = self.input
        self.artist_prompts = dict(self.prompts)
        self.artist_prompts["main"] = self.prompts["main"].format(**self.artist)
        self.agent_id = f"{self.input['ARTIST']}"

        print(f"Researching {self.input['ARTIST']}")
        print(f"Agent ID {self.agent_id}")

        mcp_url = f"http://{self.session_service.host}:{self.session_service.info['mcp_port']}/mcp"
        print(f"Using MCP session service at {mcp_url}")

        self.events = None
        self.agent_state = getattr(self, "agent_state", None)
        if self.agent_state:
            print("Resuming past agent state")
            agent = ResearchAgent.load(self.agent_state)
        else:
            print("Starting a fresh agent")
            agent = ResearchAgent(
                agent_id=self.agent_id,
                mcp_url=mcp_url,
                prompt_config=self.artist_prompts,
            )

        try:
            result = agent.run(
                self.checkpoint_interval,
                group_id=f"{self.agent_id}/{current.run_id}",
                trace_metadata={
                    "artist": self.agent_id,
                    "pathspec": current.pathspec,
                    "run-id": current.run_id,
                },
            )
        except:
            traceback.print_exc()
            self.agent_failed = getattr(self, "agent_failed", 0) + 1
            print(f"The agent has failed {self.agent_failed} times now")
            if self.agent_failed > MAX_AGENT_FAILURES:
                print("..Quit trying")
                self.decision = "failed"
            else:
                print("..Trying again")
                self.agent_state = agent.save()
                self.decision = "continue"
        else:
            if result["result"] == "timeout":
                print("Timeout - checkpointing agent state")
                self.agent_state = agent.save()
                self.decision = "continue"
            elif result["result"] == "unknown_artist":
                print(f"Can't find an artist, {self.input['ARTIST']}")
                self.decision = "unknown"
            elif result["result"] == "finished":
                self.events = result["events"]
                self.decision = "done"
            else:
                print(f"Invalid result: {result}")
                self.decision = "failed"
        self.next(
            {
                "continue": self.research,
                "done": self.done,
                "unknown": self.unknown,
                "failed": self.failed,
            },
            condition="decision",
        )

    @card(type="html")
    @step
    def done(self):
        print(f"{self.artist['ARTIST']} researched successfully!")
        values = {
            "TIMELINE_EVENTS": json.dumps(self.events),
            "ARTIST": self.artist["ARTIST"],
            "SYNOPSIS": self.artist["SYNOPSIS"],
        }
        self.html = self.timeline_html["html"]
        for tag, repl in values.items():
            self.html = self.html.replace(f"[[[{tag}]]]", repl)

        print("Events timeline:\n" + "\n".join(map(str, self.events)))
        self.next(self.agent_complete)

    @step
    def unknown(self):
        print(f"Unknown artist: {self.artist}")
        self.next(self.agent_complete)

    @step
    def failed(self):
        print(f"Research failed for {self.artist}")
        self.next(self.agent_complete)

    @step
    def agent_complete(self):
        self.next(self.join)

    @step
    def join(self, inputs):
        print("Terminate session service")
        self.session_service.terminate()
        self.next(self.end)

    @step
    def end(self):
        pass


if __name__ == "__main__":
    ResearchAgentFlow()