Agent/examples/application_reference/providers/candidates.py

from __future__ import annotations

from cyber_agent.application.candidate import CandidateAdoptionReceipt


class ReferenceCandidateRepository:
    def __init__(self, artifacts) -> None:
        self.artifacts = artifacts
        self.versions = {}
        self.adoptions = {}
        self.adoption_calls = 0

    async def put_version(self, request):
        self.versions.setdefault(request.operation_id, request)
        saved = self.versions[request.operation_id]
        return self.artifacts.put(
            artifact_id=f"candidate:{saved.candidate_id}:{saved.revision}",
            version=str(saved.revision),
            kind="candidate.output",
            root_trace_id=saved.root_trace_id,
            uid=saved.uid,
            content=saved.content,
        )

    async def merge(self, request):
        return await self.put_version(request)

    async def load_version(self, candidate_ref):
        return next(
            item.content for item in self.versions.values()
            if item.candidate_id == candidate_ref.candidate_id
            and item.revision == candidate_ref.revision
        )

    async def commit_adoption(self, request):
        self.adoption_calls += 1
        self.adoptions.setdefault(request.operation_id, request.candidate_ref)
        return CandidateAdoptionReceipt(
            operation_id=request.operation_id,
            external_id=f"published:{request.candidate_ref.candidate_id}",
        )