Save and load files with Artifacts

"""Replay the recorded Gemini turns against a real FileArtifactService in a temporary folder, offline.

Run: python -m unittest -v test_assistant.py
"""
import asyncio
import json
import logging
import tempfile
import unittest
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk import Runner  # noqa: E402
from google.adk.artifacts import FileArtifactService  # noqa: E402
from google.adk.errors.input_validation_error import InputValidationError  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.adk.sessions import InMemorySessionService  # noqa: E402
from google.genai import types  # noqa: E402

import assistant  # noqa: E402

HERE = Path(__file__).resolve().parent
RECORD = json.loads((HERE / "recorded-runs.json").read_text(encoding="utf-8"))


class ReplayLlm(BaseLlm):
    turns: list
    requests: list = []

    async def generate_content_async(self, llm_request, stream=False):
        self.requests.append(json.dumps([c.model_dump(mode="json", exclude_none=True) for c in llm_request.contents]))
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)))


class Artifacts(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.store = Path(self.tmp.name)

    def tearDown(self):
        self.tmp.cleanup()

    def replay(self, upto: int):
        """Replay steps 1..upto; a new run number means a new Runner and session on the same store."""
        async def go():
            out, runner, run = [], None, None
            for step in RECORD["steps"][:upto]:
                if step["run"] != run:
                    if runner:
                        await runner.close()
                    model = ReplayLlm(model="replay", turns=[], requests=[])
                    runner = assistant.make_runner(assistant.root_agent.clone(update={"model": model}), self.store)
                    session = await runner.session_service.create_session(app_name="docs", user_id="dana")
                    run = step["run"]
                model.turns[:] = step["model_turns"]
                start = len(model.requests)
                message = assistant.upload(HERE / "ship.toml", step["message"]) if step["kind"] == "upload" else step["message"]
                events = await assistant.ask(runner, "dana", session.id, message)
                keys = await runner.artifact_service.list_artifact_keys(app_name="docs", user_id="dana", session_id=session.id)
                out.append({"requests": model.requests[start:], "events": events, "keys": keys, "session": session.id})
            await runner.close()
            return out
        return asyncio.run(go())

    def test_an_upload_becomes_an_artifact_and_a_placeholder(self):
        step = self.replay(1)[0]
        self.assertEqual(step["keys"], ["ship.toml"])
        self.assertIn('[Uploaded Artifact: \\"ship.toml\\"]', step["requests"][0])
        self.assertNotIn("SEARCH_INDEX_URL", step["requests"][0])  # the file itself is not in the message

    def test_the_file_is_in_the_request_only_after_load_artifacts(self):
        steps = self.replay(5)
        counts = [sum("SEARCH_INDEX_URL" in r for r in s["requests"]) for s in steps]
        self.assertEqual(counts, [s["requests_with_file"] for s in RECORD["steps"]])
        self.assertEqual(counts[1:], [0, 0, 0, 0])

    def test_saving_a_filename_again_creates_a_new_version(self):
        steps = self.replay(3)
        service = FileArtifactService(root_dir=self.store)
        versions = asyncio.run(service.list_versions(app_name="docs", user_id="dana", session_id=steps[0]["session"],
                                                     filename="review.md"))
        self.assertEqual(versions, [0, 1])
        delta = {k: v for e in steps[2]["events"] for k, v in e.actions.artifact_delta.items()}
        self.assertEqual(delta, {"review.md": 1})

    def test_only_user_scoped_files_reach_a_new_session(self):
        steps = self.replay(5)
        self.assertEqual(steps[3]["keys"], ["review.md", "ship.toml", "user:checklist.md"])
        self.assertEqual(steps[4]["keys"], ["user:checklist.md"])  # review.md is still on disk, in session 1

    def test_a_parent_traversal_filename_is_refused(self):
        service = FileArtifactService(root_dir=self.store)
        with self.assertRaisesRegex(InputValidationError, "parent traversal"):
            asyncio.run(service.save_artifact(app_name="docs", user_id="dana", session_id="s",
                                              filename="../../ship.toml", artifact=types.Part(text="x")))

    def test_an_invalid_model_filename_returns_a_tool_error(self):
        model = ReplayLlm(model="replay", requests=[], turns=[
            {"role": "model", "parts": [{"function_call": {"name": "save_text_file",
             "args": {"filename": "../../outside.md", "text": "review"}, "id": "bad-name"}}]},
            {"role": "model", "parts": [{"text": "Choose a filename inside the artifact store."}]},
        ])

        async def go():
            runner = assistant.make_runner(assistant.root_agent.clone(update={"model": model}), self.store)
            session = await runner.session_service.create_session(app_name="docs", user_id="dana")
            try:
                events = await assistant.ask(runner, "dana", session.id, "Save the review.")
                keys = await runner.artifact_service.list_artifact_keys(
                    app_name="docs", user_id="dana", session_id=session.id)
                return events, keys
            finally:
                await runner.close()

        events, keys = asyncio.run(go())
        results = [r.response for event in events for r in event.get_function_responses()]
        self.assertIn("parent traversal", results[0]["error"])
        self.assertEqual(keys, [])
        self.assertTrue(events[-1].is_final_response())

    def test_a_missing_file_loads_as_none(self):
        service = FileArtifactService(root_dir=self.store)
        self.assertIsNone(asyncio.run(service.load_artifact(app_name="docs", user_id="dana", session_id="s",
                                                            filename="review.md")))

    def test_a_runner_without_an_artifact_service_fails(self):
        model = ReplayLlm(model="replay", turns=[], requests=[])
        runner = Runner(app_name="docs", agent=assistant.root_agent.clone(update={"model": model}),
                        session_service=InMemorySessionService(), auto_create_session=True)
        with self.assertRaisesRegex(ValueError, "Artifact service is not initialized"):
            asyncio.run(assistant.ask(runner, "dana", "s", "Save the text 'hello' as note.md."))
        self.assertEqual(model.requests, [])  # it fails while building the first request


if __name__ == "__main__":
    unittest.main()

"""Record two runs of the docs assistant against one artifact store on disk, on Gemini.

Run from this folder with GOOGLE_API_KEY set: python record.py
Writes recorded-runs.json and run-output.txt. The tests replay the model turns offline.
"""
import asyncio
import json
import shutil
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")

from google.adk import Runner  # noqa: E402
from google.adk.artifacts import FileArtifactService  # noqa: E402
from google.adk.sessions import InMemorySessionService  # noqa: E402
from google.genai import types  # noqa: E402

import assistant  # noqa: E402

HERE = Path(__file__).resolve().parent
STORE = HERE / "artifact-store"
STEPS = [  # (run, message): run 2 is a new process on the same store, with a new session
    (1, "upload", "Review this ship.toml before I deploy it to production."),
    (1, "text", "Save that review as review.md."),
    (1, "text", "We will fix all three problems before the deploy. Add that decision to the review and save it again."),
    (1, "text", "Save a five-item production deploy checklist for me as user:checklist.md, so I have it next time."),
    (2, "text", "Which files do you have for me? Show me my checklist."),
]


def recording_agent(turns: list, requests: list):
    def before(callback_context, llm_request):
        requests.append(json.dumps([c.model_dump(mode="json", exclude_none=True) for c in llm_request.contents]))

    def after(callback_context, llm_response):
        if llm_response.content:
            content = llm_response.content.model_dump(mode="json", exclude_none=True)
            for part in content.get("parts", []):
                part.pop("thought_signature", None)
            turns.append(content)
    return assistant.root_agent.clone(update={"before_model_callback": before, "after_model_callback": after})


def tree(store: Path) -> list[str]:
    return [str(p.relative_to(store)) for p in sorted(store.rglob("*")) if p.is_file()]


async def failures() -> dict:
    out = {}
    service = FileArtifactService(root_dir=STORE)
    # 1. A filename the model chose that walks out of the store is refused.
    for bad in ["../../ship.toml", "/etc/hosts"]:
        try:
            await service.save_artifact(app_name="docs", user_id="dana", session_id="x", filename=bad,
                                        artifact=types.Part(text="x"))
        except Exception as error:  # InputValidationError
            out.setdefault("traversal", []).append(f"{type(error).__name__}: {error}")
    # 2. Loading a file that does not exist returns None, not an error.
    out["missing"] = repr(await service.load_artifact(app_name="docs", user_id="dana", session_id="x",
                                                      filename="review.md"))
    # 3. No artifact service on the Runner: the tool's save_artifact call raises.
    runner = Runner(app_name="docs", agent=assistant.root_agent, session_service=InMemorySessionService(),
                    auto_create_session=True)
    try:
        await assistant.ask(runner, "dana", "no-store", "Save the text 'hello' as note.md.")
    except ValueError as error:
        out["no_service"] = f"{type(error).__name__}: {error}"
    await runner.close()
    return out


async def main() -> None:
    shutil.rmtree(STORE, ignore_errors=True)
    steps, runner, run = [], None, None
    for n, (run_no, kind, text) in enumerate(STEPS, 1):
        if run_no != run:
            if runner:
                await runner.close()
            turns, requests = [], []
            runner, run = assistant.make_runner(recording_agent(turns, requests)), run_no
            session = await runner.session_service.create_session(app_name="docs", user_id="dana")
        start_turns, start_requests = len(turns), len(requests)
        message = assistant.upload(HERE / "ship.toml", text) if kind == "upload" else text
        events = await assistant.ask(runner, "dana", session.id, message)
        answer = next(("".join(p.text or "" for p in e.content.parts if not p.thought).strip()
                       for e in reversed(events) if e.is_final_response() and e.content), "")
        steps.append({
            "step": n, "run": run_no, "kind": kind, "message": text,
            "calls": [{"name": c.name, "args": c.args} for e in events for c in e.get_function_calls()],
            "artifact_delta": {k: v for e in events for k, v in e.actions.artifact_delta.items()},
            "requests_with_file": sum("SEARCH_INDEX_URL" in r for r in requests[start_requests:]),
            "requests": len(requests) - start_requests,
            "model_turns": turns[start_turns:], "answer": answer,
            "listed": await runner.artifact_service.list_artifact_keys(
                app_name="docs", user_id="dana", session_id=session.id),
        })
    await runner.close()
    record = {"model": assistant.MODEL, "steps": steps, "store": tree(STORE),
              "failures": await failures()}
    (HERE / "recorded-runs.json").write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8")
    lines = []
    for s in steps:
        lines.append(f"# run {s['run']}, step {s['step']}" + (" (ship.toml attached)" if s["kind"] == "upload" else ""))
        lines.append(f"> {s['message']}")
        lines += [f"call   {c['name']} {json.dumps(c['args'])[:110]}" for c in s["calls"]]
        if s["artifact_delta"]:
            lines.append(f"saved  {s['artifact_delta']}")
        lines.append(f"listed {s['listed']}")
        lines.append("answer " + s["answer"].replace("\n", " ")[:300])
        lines.append("")
    lines.append("# artifact store on disk")
    lines += record["store"]
    (HERE / "run-output.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")
    print("\n".join(lines))
    print(json.dumps(record["failures"], indent=2))


if __name__ == "__main__":
    asyncio.run(main())