diff --git a/aperag/db/repositories/agent_runtime.py b/aperag/db/repositories/agent_runtime.py index 5eb91d2ae..2aa933f6c 100644 --- a/aperag/db/repositories/agent_runtime.py +++ b/aperag/db/repositories/agent_runtime.py @@ -18,7 +18,6 @@ from aperag.db.repositories.base import AsyncRepositoryProtocol from aperag.domains.agent_runtime.db.models import ( - AgentArtifact, AgentTimelineEvent, AgentTurn, ) @@ -182,45 +181,3 @@ async def _query(session): return result.scalars().all() return await self._execute_query(_query) - - async def create_agent_artifact( - self, - *, - turn_id: str, - artifact_type, - summary: Optional[str], - payload: dict, - storage_ref: Optional[str] = None, - ) -> AgentArtifact: - async def _operation(session): - instance = AgentArtifact( - turn_id=turn_id, - artifact_type=artifact_type, - summary=summary, - payload=payload, - storage_ref=storage_ref, - ) - session.add(instance) - await session.flush() - await session.refresh(instance) - return instance - - return await self.execute_with_transaction(_operation) - - async def query_agent_artifact(self, artifact_id: str) -> Optional[AgentArtifact]: - async def _query(session): - stmt = select(AgentArtifact).where(AgentArtifact.id == artifact_id) - result = await session.execute(stmt) - return result.scalars().first() - - return await self._execute_query(_query) - - async def query_agent_artifacts_by_turn(self, turn_id: str) -> list[AgentArtifact]: - async def _query(session): - stmt = ( - select(AgentArtifact).where(AgentArtifact.turn_id == turn_id).order_by(AgentArtifact.gmt_created.asc()) - ) - result = await session.execute(stmt) - return result.scalars().all() - - return await self._execute_query(_query) diff --git a/aperag/domains/agent_runtime/api/routes.py b/aperag/domains/agent_runtime/api/routes.py index ece47969e..2b68c7299 100644 --- a/aperag/domains/agent_runtime/api/routes.py +++ b/aperag/domains/agent_runtime/api/routes.py @@ -42,7 +42,6 @@ from aperag.domains.agent_runtime.db.models import AgentEventActor, AgentTurnStatus from aperag.domains.agent_runtime.runtime import agent_runtime_manager as runtime_manager from aperag.domains.agent_runtime.schemas import ( - AgentArtifactEnvelope, CancelTurnResponse, CreateTurnRequest, CreateTurnResponse, @@ -154,13 +153,6 @@ async def cancel_turn_view( return CancelTurnResponse(turn_id=turn_id, status=AgentTurnStatus.CANCELLED) -@router.get("/agent/artifacts/{artifact_id}") -async def get_artifact_view( - artifact_id: str, user: AuthenticatedUser = Depends(required_user) -) -> AgentArtifactEnvelope: - return await runtime_manager.artifact_service.get_artifact_for_user(str(user.id), artifact_id) - - # -- D8.3 (#75) consent + elicitation HTTP surfaces --------------- diff --git a/aperag/domains/agent_runtime/db/models.py b/aperag/domains/agent_runtime/db/models.py index 62cda39b2..8f0ae1df0 100644 --- a/aperag/domains/agent_runtime/db/models.py +++ b/aperag/domains/agent_runtime/db/models.py @@ -14,18 +14,21 @@ """Agent-runtime-domain SQLAlchemy models. -Owns ``AgentTurn`` + ``AgentTimelineEvent`` + ``AgentArtifact`` — -the three entities the agent runtime is responsible for — plus -``AgentTurnStatus`` / ``AgentEventActor`` / ``AgentArtifactType`` -enums. Moved here from ``aperag.db.models`` in Phase 5 step 5-S2b; -the legacy aggregate module retains a re-export shim until Phase 6 -cleanup. - -All three tables key on ``turn_id`` (events + artifacts) or -``chat_id`` (turn); these are opaque string FKs, so no Python-level -cross-domain import is required. ``chat_id`` is logically owned by -the ``conversation`` domain but the relationship is traversed at -query time via the turn service, not via ORM relationship() bindings. +Owns ``AgentTurn`` + ``AgentTimelineEvent`` + ``AgentMessage`` plus +``AgentTurnStatus`` / ``AgentEventActor`` enums. Moved here from +``aperag.db.models`` in Phase 5 step 5-S2b; the legacy aggregate +module retains a re-export shim until Phase 6 cleanup. + +Phase 8 D8.6 (#80) chunk-2 hard-cut removed the legacy +``AgentArtifact`` table and the ``AgentArtifactType`` enum — at-rest +answer/citation persistence is now canonical via ``AgentMessage`` +(D8.2 #74) populated by the runtime emit path. + +Tables key on ``turn_id`` (events + messages) or ``chat_id`` (turn); +these are opaque string FKs, so no Python-level cross-domain import +is required. ``chat_id`` is logically owned by the ``conversation`` +domain but the relationship is traversed at query time via the turn +service, not via ORM relationship() bindings. """ from __future__ import annotations @@ -83,14 +86,6 @@ class AgentEventActor(str, Enum): SYSTEM = "system" -class AgentArtifactType(str, Enum): - ANSWER = "answer" - REFERENCE_BUNDLE = "reference_bundle" - TOOL_RESULT_SUMMARY = "tool_result_summary" - SEARCH_RESULT_SUMMARY = "search_result_summary" - ERROR_SUMMARY = "error_summary" - - class AgentTurn(Base): __tablename__ = "agent_turn" __table_args__ = ( @@ -110,8 +105,6 @@ class AgentTurn(Base): model_profile = Column(JSON, default=lambda: {}, nullable=False) error_code = Column(String(128), nullable=True) error_message = Column(Text, nullable=True) - answer_artifact_id = Column(String(24), nullable=True, index=True) - reference_bundle_artifact_id = Column(String(24), nullable=True, index=True) timeline_cursor = Column(Integer, default=0, nullable=False) gmt_created = Column(DateTime(timezone=True), default=utc_now, nullable=False) gmt_started = Column(DateTime(timezone=True), nullable=True) @@ -138,20 +131,6 @@ class AgentTimelineEvent(Base): gmt_created = Column(DateTime(timezone=True), default=utc_now, nullable=False) -class AgentArtifact(Base): - __tablename__ = "agent_artifact" - __table_args__ = (Index("idx_agent_artifact_turn_type", "turn_id", "artifact_type"),) - - id = Column(String(24), primary_key=True, default=lambda: "art" + _random_id()) - turn_id = Column(String(24), nullable=False, index=True) - artifact_type = Column(_enum_column(AgentArtifactType), nullable=False, index=True) - summary = Column(Text, nullable=True) - payload = Column(JSON, default=lambda: {}, nullable=False) - storage_ref = Column(Text, nullable=True) - gmt_created = Column(DateTime(timezone=True), default=utc_now, nullable=False) - gmt_updated = Column(DateTime(timezone=True), default=utc_now, nullable=False) - - class AgentMessage(Base): """At-rest UIMessage envelope (Phase 8 D8.2 first-cut, D8.5-BE refined). @@ -170,10 +149,11 @@ class AgentMessage(Base): semantics independent of runtime origin per Weston msg=94dac98a / architect canonical lock msg=e01e9b4b. - The legacy ``AgentArtifact`` / ``AgentTimelineEvent`` tables are - retained alongside this table during D8.x rollout — they will be - dropped in D8.6 (#80) once D8.4 FE renderer is consuming - ``AgentMessage`` exclusively. + Phase 8 D8.6 (#80) chunk-2 dropped the legacy ``AgentArtifact`` + table and column FKs from ``AgentTurn``; ``AgentMessage`` is now + the single authoritative at-rest store for assistant turn + contents. ``AgentTimelineEvent`` remains pending chunk-3 cleanup + (replay/reload semantic change is reviewed separately). """ __tablename__ = "agent_message" diff --git a/aperag/domains/agent_runtime/runtime.py b/aperag/domains/agent_runtime/runtime.py index 0658a8d02..93018f442 100644 --- a/aperag/domains/agent_runtime/runtime.py +++ b/aperag/domains/agent_runtime/runtime.py @@ -26,8 +26,9 @@ from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.openai import OpenAIProvider +from aperag.config import get_async_session from aperag.db.ops import async_db_ops -from aperag.domains.agent_runtime.db.models import AgentArtifactType, AgentEventActor, AgentTurnStatus +from aperag.domains.agent_runtime.db.models import AgentEventActor, AgentTurnStatus from aperag.domains.agent_runtime.ports import PromptTemplateOps from aperag.domains.agent_runtime.schemas import ( AgentMessage, @@ -36,13 +37,22 @@ VisibleAgentState, ) from aperag.domains.agent_runtime.services import ( - ArtifactService, EventService, HistoryWriter, TurnService, _parse_bot_config, ) -from aperag.domains.agent_runtime.storage import DEFAULT_AGENT_TURN_LEASE_TTL_SECONDS +from aperag.domains.agent_runtime.storage import DEFAULT_AGENT_TURN_LEASE_TTL_SECONDS, AgentRuntimeRedisStore +from aperag.domains.agent_runtime.uimessage import ( + CitationData, + DataCitationPart, + SourceUrlPart, + TextPart, + UIMessage, + UIMessagePart, + UrlCitationLocation, +) +from aperag.domains.agent_runtime.uimessage_store import UIMessageDbOps, UIMessageStore from aperag.domains.knowledge_base.schemas import Collection as KBCollectionSchema from aperag.domains.model_platform.schemas import ModelCapability from aperag.exceptions import ResourceNotFoundException, ValidationException @@ -82,6 +92,43 @@ def _get_prompt_template_ops() -> PromptTemplateOps: DEFAULT_AGENT_TURN_LEASE_RENEW_INTERVAL_SECONDS = int(os.getenv("APERAG_AGENT_TURN_LEASE_RENEW_INTERVAL_SECONDS", "30")) +def _compose_assistant_parts( + *, + turn_id: str, + answer_text: str, + references: list[ReferenceBundleItem], +) -> list[UIMessagePart]: + """Project the runtime's accumulated answer + references into a + canonical at-rest ``UIMessagePart`` list for ``UIMessageStore.write``. + + Order: ``TextPart`` (answer) first, then ``SourceUrlPart`` / + ``DataCitationPart`` pairs from each reference. Items without a + URI surface as ``DataCitationPart`` only — same shape the FE + renderer expects from the live SSE stream (D8 §2 byte-equal). + """ + + parts: list[UIMessagePart] = [] + if answer_text: + parts.append(TextPart(text=answer_text)) + for index, ref in enumerate(references): + metadata = ref.metadata if isinstance(ref.metadata, dict) else {} + url = ref.uri or metadata.get("url") + title = ref.title or metadata.get("title") + snippet = ref.snippet or "" + source_id = ref.source_id or f"{turn_id}-ref-{index}" + if url: + parts.append(SourceUrlPart(source_id=str(source_id), url=str(url), title=title)) + parts.append( + DataCitationPart( + data=CitationData( + cited_text=str(snippet), + location=UrlCitationLocation(url=str(url) if url else "", title=title), + ) + ) + ) + return parts + + @dataclass class ResolvedAgentRequest: agent_message: AgentMessage @@ -181,12 +228,15 @@ async def cancel_turn(self, turn_id: str) -> None: class AgentRuntimeTaskManager: def __init__(self): self.tasks: dict[str, asyncio.Task] = {} - self.turn_service = TurnService() + self.uimessage_store = UIMessageStore( + db_ops=UIMessageDbOps(session_factory=get_async_session), + redis_store=AgentRuntimeRedisStore(), + ) + self.turn_service = TurnService(uimessage_store=self.uimessage_store) self.event_service = EventService() - self.artifact_service = ArtifactService() - self.history_writer = HistoryWriter() + self.history_writer = HistoryWriter(uimessage_store=self.uimessage_store) self.runtime: AgentRuntime = PydanticAIRuntime( - self.turn_service, self.event_service, self.artifact_service, self.history_writer + self.turn_service, self.event_service, self.history_writer, self.uimessage_store ) async def claim_turn(self, turn_id: str) -> str | None: @@ -245,13 +295,13 @@ def __init__( self, turn_service: TurnService, event_service: EventService, - artifact_service: ArtifactService, history_writer: HistoryWriter, + uimessage_store: UIMessageStore, ): self.turn_service = turn_service self.event_service = event_service - self.artifact_service = artifact_service self.history_writer = history_writer + self.uimessage_store = uimessage_store async def cancel_turn(self, turn_id: str) -> None: await self.turn_service.redis_store.mark_cancelled(turn_id) @@ -434,38 +484,19 @@ async def emit( lease_guard.ensure_owned() answer_text = "".join(text_chunks).strip() - answer_artifact = await self.artifact_service.create_artifact( + + await self.uimessage_store.write( turn_id=turn.id, - artifact_type=AgentArtifactType.ANSWER, - summary=answer_text[:200] if answer_text else "", - payload={"text": answer_text, "query": request.query}, + chat_id=chat.id, + message=UIMessage( + id=f"msg-{turn.id}", + role="assistant", + parts=_compose_assistant_parts( + turn_id=turn.id, answer_text=answer_text, references=reference_items + ), + ), ) - reference_artifact = None - if reference_items: - reference_artifact = await self.artifact_service.create_artifact( - turn_id=turn.id, - artifact_type=AgentArtifactType.REFERENCE_BUNDLE, - summary=f"{len(reference_items)} references", - payload={"items": [item.model_dump() for item in reference_items]}, - ) - await emit( - "artifact.created", - actor=AgentEventActor.SYSTEM, - label="reference_bundle", - status="ready", - data={ - "artifact_id": reference_artifact.artifact_id, - "artifact_type": reference_artifact.artifact_type, - }, - ) - await emit( - "artifact.created", - actor=AgentEventActor.SYSTEM, - label="answer", - status="ready", - data={"artifact_id": answer_artifact.artifact_id, "artifact_type": answer_artifact.artifact_type}, - ) await emit( "agent.state.changed", actor=AgentEventActor.AGENT, @@ -478,12 +509,7 @@ async def emit( label=VisibleAgentState.COMPLETED.value, status="completed", ) - await self.turn_service.mark_completed( - turn.id, - answer_artifact_id=answer_artifact.artifact_id, - reference_bundle_artifact_id=reference_artifact.artifact_id if reference_artifact else None, - sequence=sequence, - ) + await self.turn_service.mark_completed(turn.id, sequence=sequence) await self.history_writer.commit_completed_turn( turn=turn, request=request, @@ -522,19 +548,6 @@ async def emit( ) return logger.exception("Agent Runtime V3 turn failed: %s", turn.id) - error_artifact = await self.artifact_service.create_artifact( - turn_id=turn.id, - artifact_type=AgentArtifactType.ERROR_SUMMARY, - summary=str(exc), - payload={"message": str(exc), "type": exc.__class__.__name__}, - ) - await emit( - "artifact.created", - actor=AgentEventActor.SYSTEM, - label="error_summary", - status="ready", - data={"artifact_id": error_artifact.artifact_id, "artifact_type": error_artifact.artifact_type}, - ) await emit( "agent.state.changed", actor=AgentEventActor.AGENT, diff --git a/aperag/domains/agent_runtime/schemas.py b/aperag/domains/agent_runtime/schemas.py index 07497cc19..78d2fbb70 100644 --- a/aperag/domains/agent_runtime/schemas.py +++ b/aperag/domains/agent_runtime/schemas.py @@ -23,10 +23,9 @@ domain→domain. Module-level ``AGENT_RUNTIME_SCHEMA_VERSION`` constant is kept so -``AgentTurnEnvelope`` / ``AgentTimelineEventEnvelope`` / -``AgentArtifactEnvelope`` can tag every serialised payload with the -runtime contract version (declared ``v3.1``; bumping triggers -FE-side versioned parse). +``AgentTurnEnvelope`` / ``AgentTimelineEventEnvelope`` can tag every +serialised payload with the runtime contract version (declared +``v3.1``; bumping triggers FE-side versioned parse). The ``_bind_view_models_reexports`` hook at the bottom mirrors the Phase 3 step 4b / Phase 5 step 5-S3 pattern: if this module loads @@ -112,8 +111,6 @@ class AgentTurnEnvelope(BaseModel): model_profile: dict[str, Any] = Field(default_factory=dict) error_code: Optional[str] = None error_message: Optional[str] = None - answer_artifact_id: Optional[str] = None - reference_bundle_artifact_id: Optional[str] = None timeline_cursor: int = 0 started_at: Optional[datetime] = None finished_at: Optional[datetime] = None @@ -136,18 +133,6 @@ class AgentTimelineEventEnvelope(BaseModel): user_activity: Optional[UserActivityEnvelope] = None -class AgentArtifactEnvelope(BaseModel): - schema_version: str = AGENT_RUNTIME_SCHEMA_VERSION - artifact_id: str - turn_id: str - artifact_type: str - summary: Optional[str] = None - payload: dict[str, Any] = Field(default_factory=dict) - storage_ref: Optional[str] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - - class ReferenceBundleItem(BaseModel): source_type: str source_id: Optional[str] = None @@ -243,7 +228,6 @@ class AgentMessage(BaseModel): __all__ = [ "AGENT_RUNTIME_SCHEMA_VERSION", - "AgentArtifactEnvelope", "AgentMessage", "AgentTimelineEventEnvelope", "AgentTurnEnvelope", diff --git a/aperag/domains/agent_runtime/services.py b/aperag/domains/agent_runtime/services.py index 3d677d87f..b04505b89 100644 --- a/aperag/domains/agent_runtime/services.py +++ b/aperag/domains/agent_runtime/services.py @@ -21,7 +21,6 @@ from aperag.db.ops import AsyncDatabaseOps, async_db_ops from aperag.domains.agent_runtime.db.models import AgentEventActor, AgentTurnStatus from aperag.domains.agent_runtime.schemas import ( - AgentArtifactEnvelope, AgentTimelineEventEnvelope, AgentTurnEnvelope, CreateTurnRequest, @@ -30,12 +29,8 @@ UserActivityEnvelope, UserActivityIntent, ) -from aperag.domains.agent_runtime.snapshot_assembler import ( - assemble_parts_from_artifacts, - extract_error_text, -) from aperag.domains.agent_runtime.storage import AgentRuntimeRedisStore -from aperag.domains.agent_runtime.uimessage import AgentTurnSnapshot +from aperag.domains.agent_runtime.uimessage import AgentTurnSnapshot, TextPart from aperag.domains.agent_runtime.uimessage_store import UIMessageStore from aperag.domains.conversation.db.models import BotType from aperag.domains.conversation.schemas import BotConfig @@ -67,10 +62,6 @@ def _apply_runtime_state_to_turn( updates["timeline_cursor"] = max(timeline_cursor, _coerce_timeline_cursor(runtime_state.get("timeline_cursor"))) if "status" in runtime_state: updates["status"] = runtime_state["status"] - if runtime_state.get("answer_artifact_id"): - updates["answer_artifact_id"] = runtime_state["answer_artifact_id"] - if runtime_state.get("reference_bundle_artifact_id"): - updates["reference_bundle_artifact_id"] = runtime_state["reference_bundle_artifact_id"] if runtime_state.get("error_code"): updates["error_code"] = runtime_state["error_code"] if runtime_state.get("error_message"): @@ -292,12 +283,6 @@ def _build_user_activity_for_event( return _build_user_activity(UserActivityIntent.WAITING) -def _extract_answer_text_from_artifact(artifact) -> str: - if not artifact or not isinstance(artifact.payload, dict): - return "" - return artifact.payload.get("text") or artifact.payload.get("content") or "" - - class TurnService: def __init__( self, @@ -307,11 +292,6 @@ def __init__( ): self.db_ops = db_ops or async_db_ops self.redis_store = redis_store or AgentRuntimeRedisStore() - # ``uimessage_store`` is the canonical at-rest reader once the - # wire emitter starts persisting ``agent_message.parts`` (D8.6 - # / #80). Until then it stays optional and we fall back to - # projecting legacy artifacts via - # ``assemble_parts_from_artifacts``. self.uimessage_store = uimessage_store async def get_chat_and_bot(self, user: str, chat_id: str): @@ -364,19 +344,10 @@ async def mark_running(self, turn_id: str) -> None: {"status": turn.status, "timeline_cursor": turn.timeline_cursor, "chat_id": turn.chat_id}, ) - async def mark_completed( - self, - turn_id: str, - *, - answer_artifact_id: Optional[str], - reference_bundle_artifact_id: Optional[str], - sequence: int, - ) -> None: + async def mark_completed(self, turn_id: str, *, sequence: int) -> None: turn = await self.db_ops.update_agent_turn( turn_id, status=AgentTurnStatus.COMPLETED, - answer_artifact_id=answer_artifact_id, - reference_bundle_artifact_id=reference_bundle_artifact_id, timeline_cursor=sequence, gmt_finished=utc_now(), ) @@ -387,8 +358,6 @@ async def mark_completed( "status": turn.status, "timeline_cursor": turn.timeline_cursor, "chat_id": turn.chat_id, - "answer_artifact_id": answer_artifact_id, - "reference_bundle_artifact_id": reference_bundle_artifact_id, }, ) @@ -427,28 +396,15 @@ async def mark_cancelled(self, turn_id: str, *, sequence: int) -> None: ) async def get_turn_snapshot(self, user: str, chat_id: str, turn_id: str) -> AgentTurnSnapshot: - """Return the canonical ``AgentTurnSnapshot`` (Phase 8 D8.4d). - - The legacy ``{turn, timeline, artifacts}`` shape is gone; the - FE renderer (#76 / #77 / #78) now consumes the same - ``UIMessagePart`` discriminated union it already gets from the - live SSE stream (per D8 §2 wire / at-rest byte-equal). - - Sources, in priority order: - - 1. ``UIMessageStore.read(turn_id)`` — once the wire emitter - starts persisting ``agent_message.parts`` (D8.6 / #80 - cleanup) this is the only path. Today the store is - optional and most reads return ``None``. - 2. ``assemble_parts_from_artifacts`` — transitional projection - from the legacy ``answer`` / ``reference_bundle`` - artifacts so completed / failed / cancelled turns keep - rendering before D8.6. - - The ``error_text`` field is filled from the - ``error_summary`` artifact (or the runtime ``error_message`` - field as a fallback) so a FAILED reload shows the same - message the live ``error`` part would have surfaced. + """Return the canonical ``AgentTurnSnapshot`` from the at-rest UIMessage store. + + Phase 8 D8.6 (#80) chunk-2 hard-cut removed the legacy + ``AgentArtifact`` projection fallback. The live runtime now + writes ``UIMessageStore`` at end-of-turn (D8.2 #74 store + + D8.6 #80 wire-in), so the snapshot endpoint reads the + canonical ``UIMessage`` directly. ``error_text`` falls back + to the ``AgentTurn`` row's ``error_message`` for FAILED / + CANCELLED turns. """ turn = await self.db_ops.query_agent_turn(user, chat_id, turn_id) @@ -464,27 +420,14 @@ async def get_turn_snapshot(self, user: str, chat_id: str, turn_id: str) -> Agen ) parts: list[Any] = [] - artifacts: list[Any] = [] - if self.uimessage_store is not None: - try: - persisted_message = await self.uimessage_store.read(turn_id) - except Exception: # pragma: no cover — store is best-effort during transition - persisted_message = None + persisted_message = await self.uimessage_store.read(turn_id) if persisted_message is not None and persisted_message.parts: parts = list(persisted_message.parts) - if not parts: - artifacts = await self.db_ops.query_agent_artifacts_by_turn(turn_id) - parts = list(assemble_parts_from_artifacts(artifacts)) - error_text: Optional[str] = None if status_str in {AgentTurnStatus.FAILED.value, AgentTurnStatus.CANCELLED.value}: - if not artifacts: - artifacts = await self.db_ops.query_agent_artifacts_by_turn(turn_id) - error_text = extract_error_text(artifacts) - if not error_text: - error_text = runtime_state.get("error_message") or turn.error_message + error_text = runtime_state.get("error_message") or turn.error_message return AgentTurnSnapshot( turn_id=turn.id, @@ -521,8 +464,6 @@ def to_turn_envelope(turn) -> AgentTurnEnvelope: model_profile=turn.model_profile or {}, error_code=turn.error_code, error_message=turn.error_message, - answer_artifact_id=turn.answer_artifact_id, - reference_bundle_artifact_id=turn.reference_bundle_artifact_id, timeline_cursor=turn.timeline_cursor or 0, started_at=turn.gmt_started, finished_at=turn.gmt_finished, @@ -605,91 +546,38 @@ def adapt_event_envelope(event: AgentTimelineEventEnvelope) -> AgentTimelineEven ) -class ArtifactService: - def __init__(self, db_ops: AsyncDatabaseOps | None = None): - self.db_ops = db_ops or async_db_ops - - async def create_artifact( +class HistoryWriter: + def __init__( self, - *, - turn_id: str, - artifact_type, - summary: Optional[str], - payload: dict[str, Any], - storage_ref: Optional[str] = None, - ) -> AgentArtifactEnvelope: - artifact = await self.db_ops.create_agent_artifact( - turn_id=turn_id, - artifact_type=artifact_type, - summary=summary, - payload=payload, - storage_ref=storage_ref, - ) - return self.to_artifact_envelope(artifact) - - async def get_artifact_for_user(self, user: str, artifact_id: str) -> AgentArtifactEnvelope: - artifact = await self.db_ops.query_agent_artifact(artifact_id) - if not artifact: - raise ResourceNotFoundException("Artifact", artifact_id) - - turn = await self._query_turn_for_user(user, artifact.turn_id) - if not turn: - raise ResourceNotFoundException("Artifact", artifact_id) - return self.to_artifact_envelope(artifact) - - async def _query_turn_for_user(self, user: str, turn_id: str): - async def _query(session): - from sqlalchemy import select - - from aperag.domains.agent_runtime.db.models import AgentTurn - - stmt = select(AgentTurn).where(AgentTurn.id == turn_id, AgentTurn.user == user) - result = await session.execute(stmt) - return result.scalars().first() - - return await self.db_ops._execute_query(_query) + db_ops: AsyncDatabaseOps | None = None, + uimessage_store: UIMessageStore | None = None, + ): + self.db_ops = db_ops or async_db_ops + self.uimessage_store = uimessage_store - @staticmethod - def to_artifact_envelope(artifact) -> AgentArtifactEnvelope: - artifact_type = ( - artifact.artifact_type.value if hasattr(artifact.artifact_type, "value") else artifact.artifact_type - ) - return AgentArtifactEnvelope( - artifact_id=artifact.id, - turn_id=artifact.turn_id, - artifact_type=artifact_type, - summary=artifact.summary, - payload=artifact.payload or {}, - storage_ref=artifact.storage_ref, - created_at=artifact.gmt_created, - updated_at=artifact.gmt_updated, - ) + async def build_history_context(self, user: str, chat_id: str, limit: int = 8) -> str: + """Compose a plain-text history of recent COMPLETED turns for prompt context. + Phase 8 D8.6 (#80) chunk-2 hard-cut switched the source from + the legacy ``AgentArtifact`` rows to the canonical + ``UIMessage`` parts written by the runtime at end-of-turn. + Each turn contributes its persisted assistant ``TextPart`` + text (joined when multiple) plus the original ``input_text``. + """ -class HistoryWriter: - def __init__(self, db_ops: AsyncDatabaseOps | None = None): - self.db_ops = db_ops or async_db_ops + if self.uimessage_store is None: + return "" - async def build_history_context(self, user: str, chat_id: str, limit: int = 8) -> str: turns = await self.db_ops.query_recent_agent_turns(user, chat_id, limit=limit) lines: list[str] = [] for turn in turns: if turn.status != AgentTurnStatus.COMPLETED: continue - artifacts = await self.db_ops.query_agent_artifacts_by_turn(turn.id) - answer = None - if turn.answer_artifact_id: - answer = next((artifact for artifact in artifacts if artifact.id == turn.answer_artifact_id), None) - if not answer: - answer = next( - ( - artifact - for artifact in artifacts - if getattr(artifact.artifact_type, "value", artifact.artifact_type) == "answer" - ), - None, - ) - answer_text = _extract_answer_text_from_artifact(answer) + message = await self.uimessage_store.read(turn.id) + if message is None or not message.parts: + continue + chunks = [part.text for part in message.parts if isinstance(part, TextPart) and part.text] + answer_text = "\n".join(chunks).strip() if not answer_text: continue lines.append(f"User: {turn.input_text}") diff --git a/aperag/domains/agent_runtime/snapshot_assembler.py b/aperag/domains/agent_runtime/snapshot_assembler.py deleted file mode 100644 index ee5b64015..000000000 --- a/aperag/domains/agent_runtime/snapshot_assembler.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright 2025 ApeCloud, Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Snapshot endpoint UIMessage assembler (Phase 8 D8.4d / #90). - -The agent runtime's snapshot endpoint historically returned the legacy -``{turn, timeline, artifacts}`` shape. D8 §2 locks the wire and at-rest -``UIMessage`` parts as same-schema, so the snapshot endpoint must -expose ``parts: UIMessagePart[]`` for the FE renderer (#76 / #77 / #78) -to consume from a single code path. - -This module projects the legacy ``AgentArtifact`` rows into the -canonical at-rest part list, mirroring the FE-side -``snapshot-fallback.ts`` adapter that #77 huangheng landed as a -transitional bridge: - -* ``answer`` artifact -> single ``TextPart`` -* ``reference_bundle`` artifact -> N x ``SourceUrlPart`` plus - N x ``DataCitationPart`` (Anthropic-shape) -* ``error_summary`` artifact -> not a part; surfaced via the snapshot - envelope's ``error_text`` field - -Other artifact types (``tool_result_summary`` / -``search_result_summary``) are skipped: tool lifecycle and search -results are surfaced by the tool/citation parts emitted directly during -the live turn (#73 wire emitter + #75 tool lifecycle), so projecting -them here would double-count once the wire emitter starts persisting -``agent_message.parts`` directly (D8.6 / #80 cleanup). - -When that cleanup lands and ``UIMessageStore.read(turn_id)`` returns a -populated row, the snapshot service short-circuits to the at-rest read -and this assembler is no longer hit. At that point the FE -``snapshot-fallback.ts`` and this module can both be deleted in a -single follow-up. -""" - -from __future__ import annotations - -from typing import Iterable, Optional - -from aperag.domains.agent_runtime.db.models import AgentArtifact, AgentArtifactType -from aperag.domains.agent_runtime.uimessage import ( - CitationData, - DataCitationPart, - SourceUrlPart, - TextPart, - UIMessagePart, - UrlCitationLocation, -) - - -def assemble_parts_from_artifacts( - artifacts: Iterable[AgentArtifact], -) -> list[UIMessagePart]: - """Project legacy ``AgentArtifact`` rows into a UIMessage parts list. - - Order: ``TextPart`` (answer) first, then ``SourceUrlPart`` / - ``DataCitationPart`` pairs from ``reference_bundle.items``. The - answer text comes from ``payload.text`` (matches the runtime's - artifact write at ``runtime.py:441``); reference items are read - from ``payload.items`` (matches ``runtime.py:449``). - - Items that lack a URI are still surfaced as ``DataCitationPart`` - with an empty ``url`` so the FE renderer can show snippet/title; - the upcoming D8.6 cleanup will replace this transitional adapter - with the canonical at-rest ``UIMessageStore.read()`` path. - """ - - answer: Optional[AgentArtifact] = None - reference_bundle: Optional[AgentArtifact] = None - for artifact in artifacts: - if artifact.artifact_type == AgentArtifactType.ANSWER and answer is None: - answer = artifact - elif artifact.artifact_type == AgentArtifactType.REFERENCE_BUNDLE and reference_bundle is None: - reference_bundle = artifact - - parts: list[UIMessagePart] = [] - - if answer is not None: - text = _extract_answer_text(answer) - if text: - parts.append(TextPart(text=text)) - - if reference_bundle is not None: - for source_part, citation_part in _project_reference_items(reference_bundle): - if source_part is not None: - parts.append(source_part) - parts.append(citation_part) - - return parts - - -def extract_error_text(artifacts: Iterable[AgentArtifact]) -> Optional[str]: - """Return the human-readable error message for a failed turn, or None. - - Looks up the ``error_summary`` artifact and pulls - ``payload.message`` / ``payload.text`` / ``summary`` in that order. - Mirrors the FE adapter's ``extractErrorTextFromSnapshot`` so a - historical FAILED turn renders the same error string from either - side of the new boundary. - """ - - for artifact in artifacts: - if artifact.artifact_type != AgentArtifactType.ERROR_SUMMARY: - continue - payload = artifact.payload if isinstance(artifact.payload, dict) else {} - message = payload.get("message") or payload.get("text") or payload.get("summary") - if message: - return str(message) - if artifact.summary: - return str(artifact.summary) - return None - - -def _extract_answer_text(artifact: AgentArtifact) -> str: - payload = artifact.payload if isinstance(artifact.payload, dict) else {} - text = payload.get("text") or payload.get("content") or payload.get("answer") - if text: - return str(text) - if artifact.summary: - return str(artifact.summary) - return "" - - -def _project_reference_items( - artifact: AgentArtifact, -) -> Iterable[tuple[Optional[SourceUrlPart], DataCitationPart]]: - payload = artifact.payload if isinstance(artifact.payload, dict) else {} - raw_items = payload.get("items") - if not isinstance(raw_items, list): - return - - for index, raw in enumerate(raw_items): - if not isinstance(raw, dict): - continue - metadata = raw.get("metadata") if isinstance(raw.get("metadata"), dict) else {} - url = raw.get("uri") or metadata.get("url") - title = raw.get("title") or metadata.get("title") - snippet = raw.get("snippet") or raw.get("content") or "" - source_id = raw.get("source_id") or raw.get("id") or f"{artifact.id}-ref-{index}" - - source_part: Optional[SourceUrlPart] = None - if url: - source_part = SourceUrlPart(source_id=str(source_id), url=str(url), title=title) - - citation_part = DataCitationPart( - data=CitationData( - cited_text=str(snippet), - location=UrlCitationLocation(url=str(url) if url else "", title=title), - ) - ) - yield source_part, citation_part diff --git a/aperag/domains/conversation/service/chat_completion_service.py b/aperag/domains/conversation/service/chat_completion_service.py index 9bb4adb57..5752050f9 100644 --- a/aperag/domains/conversation/service/chat_completion_service.py +++ b/aperag/domains/conversation/service/chat_completion_service.py @@ -366,14 +366,13 @@ async def _collect_openai_turn_response( if self._status_value(turn.status) != AgentTurnStatus.COMPLETED.value: return OpenAIFormatter.format_error(turn.error_message or "Chat completion failed") - # Phase 8 D8.4d (#90): the snapshot endpoint now returns - # canonical UIMessage parts; legacy ``snapshot.artifacts`` - # is gone. Pull artifacts directly from the DB so the - # OpenAI-compat completion content keeps its existing - # answer + references shape (which is independent of the - # FE-facing UIMessage protocol). - artifacts = await runtime_manager.turn_service.db_ops.query_agent_artifacts_by_turn(turn_id) - content = self._build_completion_content(artifacts) + # Phase 8 D8.6 (#80) chunk-2 hard-cut: pull the canonical + # UIMessage parts persisted by the runtime at end-of-turn + # and project them into the OpenAI-compat completion + # content shape (answer text + references payload). + persisted = await runtime_manager.uimessage_store.read(turn_id) + parts = list(persisted.parts) if persisted and persisted.parts else [] + content = self._build_completion_content(parts) return OpenAIFormatter.format_complete_response(turn_id, content) finally: if ephemeral_chat: @@ -444,30 +443,45 @@ def _status_value(status: Any) -> str: return status.value if hasattr(status, "value") else str(status) @staticmethod - def _build_completion_content(artifacts) -> str: - """Build OpenAI-compat completion content from raw agent artifacts. - - Phase 8 D8.4d (#90) flipped the snapshot endpoint to canonical - UIMessage parts; this OpenAI-compat path stays artifact-shaped - because the response format is independent of the FE-facing - UIMessage protocol. Callers pass the raw artifact list from - ``db_ops.query_agent_artifacts_by_turn``. + def _build_completion_content(parts) -> str: + """Build OpenAI-compat completion content from canonical UIMessage parts. + + Phase 8 D8.6 (#80) chunk-2 switched the source from legacy + ``AgentArtifact`` rows to the at-rest ``UIMessage`` parts the + runtime now writes at end-of-turn. The answer text comes from + joined ``TextPart`` content; the references payload is + derived from ``DataCitationPart`` entries (one dict per + citation, preserving ``url`` / ``title`` / ``cited_text`` so + the OpenAI-compat envelope keeps the existing answer + + references shape). """ - answer_text = "" + from aperag.domains.agent_runtime.uimessage import ( + DataCitationPart, + TextPart, + UrlCitationLocation, + ) + + text_chunks: list[str] = [] references_payload: list[dict[str, Any]] = [] - for artifact in artifacts: - artifact_type = getattr(artifact, "artifact_type", None) - type_value = getattr(artifact_type, "value", artifact_type) - payload = getattr(artifact, "payload", None) or {} - if type_value == "answer": - if not answer_text: - answer_text = payload.get("text") or payload.get("content") or "" - elif type_value == "reference_bundle": - items = payload.get("items") - if isinstance(items, list): - references_payload = [item for item in items if isinstance(item, dict)] + for part in parts: + if isinstance(part, TextPart) and part.text: + text_chunks.append(part.text) + continue + if isinstance(part, DataCitationPart): + location = part.data.location if part.data else None + url = location.url if isinstance(location, UrlCitationLocation) else None + title = location.title if isinstance(location, UrlCitationLocation) else None + cited_text = part.data.cited_text if part.data else "" + references_payload.append( + { + "uri": url or "", + "title": title, + "snippet": cited_text, + } + ) + answer_text = "".join(text_chunks) if references_payload: return f"{answer_text}{DOC_QA_REFERENCES}{json.dumps(references_payload, ensure_ascii=False)}" return answer_text diff --git a/aperag/domains/conversation/service/chat_service.py b/aperag/domains/conversation/service/chat_service.py index ad43eb725..f105d121e 100644 --- a/aperag/domains/conversation/service/chat_service.py +++ b/aperag/domains/conversation/service/chat_service.py @@ -23,8 +23,7 @@ * ``Chat`` / ``ChatStatus`` → ``aperag.domains.conversation.db.models`` * ``BotType`` → ``aperag.domains.conversation.db.models`` -* ``AgentArtifactType`` / ``AgentTurnStatus`` → - ``aperag.domains.agent_runtime.db.models`` +* ``AgentTurnStatus`` → ``aperag.domains.agent_runtime.db.models`` * Pydantic ``Chat`` / ``ChatDetails`` / ``ChatMessage`` / ``Reference`` → ``aperag.domains.conversation.schemas`` @@ -41,11 +40,9 @@ from aperag.db.ops import AsyncDatabaseOps, async_db_ops from aperag.domains.agent_runtime.db.models import AgentTurnStatus -from aperag.domains.agent_runtime.snapshot_assembler import ( - assemble_parts_from_artifacts, - extract_error_text, -) +from aperag.domains.agent_runtime.storage import AgentRuntimeRedisStore from aperag.domains.agent_runtime.uimessage import AgentTurnSnapshot +from aperag.domains.agent_runtime.uimessage_store import UIMessageDbOps, UIMessageStore from aperag.domains.conversation.db.models import BotType, ChatStatus from aperag.domains.conversation.db.models import Chat as ChatRow from aperag.domains.conversation.schemas import Chat, ChatDetails, ChatUpdate @@ -57,11 +54,28 @@ class ChatService: """Chat service that handles business logic for chats""" - def __init__(self, session: AsyncSession = None): + def __init__( + self, + session: AsyncSession = None, + *, + uimessage_store: Optional[UIMessageStore] = None, + ): if session is None: self.db_ops = async_db_ops else: self.db_ops = AsyncDatabaseOps(session) + self.uimessage_store = uimessage_store + + async def _resolve_uimessage_store(self) -> UIMessageStore: + if self.uimessage_store is not None: + return self.uimessage_store + from aperag.config import get_async_session + + self.uimessage_store = UIMessageStore( + db_ops=UIMessageDbOps(session_factory=get_async_session), + redis_store=AgentRuntimeRedisStore(), + ) + return self.uimessage_store def build_chat_response(self, chat: ChatRow) -> Chat: """Build Chat response object for API return.""" @@ -78,45 +92,24 @@ def build_chat_response(self, chat: ChatRow) -> Chat: async def _build_v3_chat_history(self, user: str, chat_id: str) -> list[AgentTurnSnapshot]: """Build the chat history as canonical ``AgentTurnSnapshot`` envelopes. - Phase 8 D8.5-BE (#92) flips this from the legacy - ``list[list[ChatMessage]]`` shape to one ``AgentTurnSnapshot`` - per assistant turn. Each snapshot carries the same - ``UIMessagePart[]`` shape the FE consumes from the live SSE - stream, so historical and live turns render through a single - canonical path (D8 §2 wire/at-rest byte-equal). - - The user query is exposed at ``input_text`` on the snapshot - envelope rather than as a separate ``role=human`` ChatMessage - entry; the FE renders it from there. The assistant turn's - ``answer`` / ``reference_bundle`` / ``error_summary`` - artifacts are projected into ``parts`` via - :func:`assemble_parts_from_artifacts`, mirroring the snapshot - endpoint (#90 / D8.4d). FAILED / CANCELLED turns surface their - message via ``error_text`` (preferring an ``error_summary`` - artifact, falling back to ``turn.error_message``), again - mirroring the snapshot endpoint contract. - - Once the wire emitter starts populating ``agent_message.parts`` - directly (D8.6 / #80), this method can short-circuit to - :meth:`UIMessageStore.read` per-turn; until then the artifact - projection is the single source. ``runtime_kind`` is hardcoded - to ``agent_runtime`` here — non-agent runtimes will write - ``direct_chat`` / ``rag_chat`` rows directly via the future - non-agent write path and this method only needs to surface - them when they exist (a no-op until then per architect lock - msg=01918929). + Phase 8 D8.6 (#80) chunk-2 hard-cut switched the per-turn + ``parts`` source from the legacy ``AgentArtifact`` projection + to the canonical ``UIMessage`` rows the runtime now writes at + end-of-turn. FAILED / CANCELLED turns surface their message + via ``error_text`` from the ``AgentTurn`` row directly. """ turns = await self.db_ops.query_agent_turns(user, chat_id) + store = await self._resolve_uimessage_store() history: list[AgentTurnSnapshot] = [] for turn in turns: - artifacts = await self.db_ops.query_agent_artifacts_by_turn(turn.id) - parts = list(assemble_parts_from_artifacts(artifacts)) + persisted_message = await store.read(turn.id) + parts = list(persisted_message.parts) if persisted_message and persisted_message.parts else [] error_text: Optional[str] = None if turn.status in {AgentTurnStatus.FAILED, AgentTurnStatus.CANCELLED}: - error_text = extract_error_text(artifacts) or turn.error_message + error_text = turn.error_message status_value = turn.status.value if hasattr(turn.status, "value") else str(turn.status) history.append( diff --git a/aperag/domains/evaluation/worker.py b/aperag/domains/evaluation/worker.py index 2a3e82a05..a35771ba3 100644 --- a/aperag/domains/evaluation/worker.py +++ b/aperag/domains/evaluation/worker.py @@ -172,12 +172,12 @@ async def dispatch_evaluation_turn( latency_ms = int((time.monotonic() - start) * 1000) if final_status == AgentTurnStatus.COMPLETED.value: - # Phase 8 D8.4d (#90): the snapshot endpoint now returns canonical - # UIMessage parts; legacy ``snapshot.artifacts`` is gone. Pull - # artifacts straight from the DB instead of rebuilding the answer - # from the at-rest UIMessage parts (which is the FE concern). - artifacts = await agent_runtime_manager.turn_service.db_ops.query_agent_artifacts_by_turn(turn.id) - answer_text = _extract_answer_text(artifacts) + # Phase 8 D8.6 (#80) chunk-2: read canonical UIMessage parts + # the runtime persists at end-of-turn and join all TextPart + # contributions into the evaluation answer text. + persisted = await agent_runtime_manager.uimessage_store.read(turn.id) + parts = list(persisted.parts) if persisted and persisted.parts else [] + answer_text = _extract_answer_text(parts) return TurnDispatchOutcome( status=EvaluationRunItemAttemptStatus.COMPLETED, agent_chat_id=chat_id, @@ -211,18 +211,12 @@ async def dispatch_evaluation_turn( ) -def _extract_answer_text(artifacts) -> str: - for artifact in artifacts: - artifact_type = getattr(artifact, "artifact_type", None) - # AgentArtifactType is a str-Enum; compare by string value to - # avoid pulling the enum into this module just for one branch. - type_value = getattr(artifact_type, "value", artifact_type) - if type_value == "answer": - payload = getattr(artifact, "payload", None) or {} - text = payload.get("text") or payload.get("content") - if text: - return text - return "" +def _extract_answer_text(parts) -> str: + """Join the assistant's ``TextPart`` contents into a single string.""" + from aperag.domains.agent_runtime.uimessage import TextPart + + chunks = [part.text for part in parts if isinstance(part, TextPart) and part.text] + return "".join(chunks) # --------------------------------------------------------------------------- diff --git a/aperag/migration/versions/20260426020000-d8e6c2b4f1a9_drop_agent_artifact.py b/aperag/migration/versions/20260426020000-d8e6c2b4f1a9_drop_agent_artifact.py new file mode 100644 index 000000000..c06d9a09a --- /dev/null +++ b/aperag/migration/versions/20260426020000-d8e6c2b4f1a9_drop_agent_artifact.py @@ -0,0 +1,71 @@ +"""drop agent_artifact table + agent_turn artifact_id columns (D8.6 #80 chunk-2) + +Phase 8 D8.6 (#80) chunk-2 hard-cut: the agent runtime now writes the +canonical ``UIMessage`` envelope to ``agent_message`` at end-of-turn +(D8.2 #74 store + #80 wire-in), so the legacy ``agent_artifact`` row +projection (``answer`` / ``reference_bundle`` / ``error_summary`` / +``tool_result_summary`` / ``search_result_summary``) is dead weight. +This migration drops the table and the two ``agent_turn`` FK columns +that pointed into it. + +Pre-launch system has no users / no data, so the cutover is direct +delete (per earayu2 hard-cut acceptance) — no backfill / no row +migration. The downgrade restores the dropped surface so a rollback +can replay subsequent migrations cleanly. + +The ``agent_timeline_event`` table is intentionally retained — its +removal is chunk-3 (replay/reload semantic change reviewed +separately). + +Revision ID: d8e6c2b4f1a9 +Revises: c8f2d34a51e7 +Create Date: 2026-04-26 02:00:00.000000 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "d8e6c2b4f1a9" +down_revision: Union[str, None] = "c8f2d34a51e7" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.drop_index("ix_agent_turn_reference_bundle_artifact_id", table_name="agent_turn") + op.drop_index("ix_agent_turn_answer_artifact_id", table_name="agent_turn") + op.drop_column("agent_turn", "reference_bundle_artifact_id") + op.drop_column("agent_turn", "answer_artifact_id") + op.drop_index("ix_agent_artifact_turn_id", table_name="agent_artifact") + op.drop_index("ix_agent_artifact_artifact_type", table_name="agent_artifact") + op.drop_index("idx_agent_artifact_turn_type", table_name="agent_artifact") + op.drop_table("agent_artifact") + + +def downgrade() -> None: + op.create_table( + "agent_artifact", + sa.Column("id", sa.String(length=24), nullable=False), + sa.Column("turn_id", sa.String(length=24), nullable=False), + sa.Column("artifact_type", sa.String(length=50), nullable=False), + sa.Column("summary", sa.Text(), nullable=True), + sa.Column("payload", sa.JSON(), nullable=False), + sa.Column("storage_ref", sa.Text(), nullable=True), + sa.Column("gmt_created", sa.DateTime(timezone=True), nullable=False), + sa.Column("gmt_updated", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("idx_agent_artifact_turn_type", "agent_artifact", ["turn_id", "artifact_type"], unique=False) + op.create_index("ix_agent_artifact_artifact_type", "agent_artifact", ["artifact_type"], unique=False) + op.create_index("ix_agent_artifact_turn_id", "agent_artifact", ["turn_id"], unique=False) + op.add_column("agent_turn", sa.Column("answer_artifact_id", sa.String(length=24), nullable=True)) + op.add_column("agent_turn", sa.Column("reference_bundle_artifact_id", sa.String(length=24), nullable=True)) + op.create_index("ix_agent_turn_answer_artifact_id", "agent_turn", ["answer_artifact_id"], unique=False) + op.create_index( + "ix_agent_turn_reference_bundle_artifact_id", + "agent_turn", + ["reference_bundle_artifact_id"], + unique=False, + ) diff --git a/tests/unit_test/agent_runtime/test_agent_runtime_openapi_contract.py b/tests/unit_test/agent_runtime/test_agent_runtime_openapi_contract.py index 524e2494f..ce1558eef 100644 --- a/tests/unit_test/agent_runtime/test_agent_runtime_openapi_contract.py +++ b/tests/unit_test/agent_runtime/test_agent_runtime_openapi_contract.py @@ -9,7 +9,13 @@ def _json_schema(operation: dict, status: str = "200") -> dict: return operation["responses"][status]["content"]["application/json"]["schema"] -def test_agent_runtime_v2_openapi_contract_exposes_turn_event_and_artifact_models(): +def test_agent_runtime_v2_openapi_contract_exposes_turn_event_models(): + """Phase 8 D8.6 (#80) chunk-2 dropped the legacy + ``/api/v2/agent/artifacts/{artifact_id}`` endpoint and the + ``AgentArtifactEnvelope`` schema; the OpenAPI surface should no + longer expose either. + """ + create_turn = _operation("/api/v2/agent/chats/{chat_id}/turns", "post") assert create_turn["requestBody"]["content"]["application/json"]["schema"] == { "$ref": "#/components/schemas/CreateTurnRequest" @@ -22,8 +28,10 @@ def test_agent_runtime_v2_openapi_contract_exposes_turn_event_and_artifact_model cancel = _operation("/api/v2/agent/chats/{chat_id}/turns/{turn_id}/cancel", "post") assert _json_schema(cancel) == {"$ref": "#/components/schemas/CancelTurnResponse"} - artifact = _operation("/api/v2/agent/artifacts/{artifact_id}", "get") - assert _json_schema(artifact) == {"$ref": "#/components/schemas/AgentArtifactEnvelope"} + paths = app.openapi()["paths"] + assert "/api/v2/agent/artifacts/{artifact_id}" not in paths + schemas = app.openapi()["components"]["schemas"] + assert "AgentArtifactEnvelope" not in schemas events = _operation("/api/v2/agent/chats/{chat_id}/turns/{turn_id}/events", "get") event_stream = events["responses"]["200"]["content"]["text/event-stream"]["schema"] diff --git a/tests/unit_test/agent_runtime/test_agent_runtime_v3.py b/tests/unit_test/agent_runtime/test_agent_runtime_v3.py index 0edd9d0a3..8d0a562c6 100644 --- a/tests/unit_test/agent_runtime/test_agent_runtime_v3.py +++ b/tests/unit_test/agent_runtime/test_agent_runtime_v3.py @@ -6,19 +6,22 @@ import aperag.domains.agent_runtime.storage as agent_runtime_storage from aperag.domains.agent_runtime.api import routes as agent_runtime_view -from aperag.domains.agent_runtime.db.models import ( - AgentArtifactType, - AgentEventActor, - AgentTurnStatus, -) +from aperag.domains.agent_runtime.db.models import AgentEventActor, AgentTurnStatus from aperag.domains.agent_runtime.schemas import ( - AgentArtifactEnvelope, AgentTimelineEventEnvelope, CreateTurnRequest, UserActivityIntent, ) from aperag.domains.agent_runtime.services import EventService, HistoryWriter, TurnService -from aperag.domains.agent_runtime.uimessage import AgentTurnSnapshot +from aperag.domains.agent_runtime.uimessage import ( + AgentTurnSnapshot, + CitationData, + DataCitationPart, + SourceUrlPart, + TextPart, + UIMessage, + UrlCitationLocation, +) def _now(): @@ -70,11 +73,20 @@ async def eval(self, script, _numkeys, key, token, *_args): raise AssertionError(f"Unexpected script: {script}") +class _FakeUIMessageStore: + """Minimal in-memory UIMessageStore stand-in for unit tests.""" + + def __init__(self, *, messages=None): + self.messages = dict(messages or {}) + + async def read(self, turn_id): + return self.messages.get(turn_id) + + class _FakeDbOps: - def __init__(self, *, turn=None, persisted_events=None, artifacts=None): + def __init__(self, *, turn=None, persisted_events=None): self.turn = turn self.persisted_events = list(persisted_events or []) - self.artifacts = list(artifacts or []) async def query_agent_turn(self, _user, _chat_id, _turn_id): return self.turn @@ -82,9 +94,6 @@ async def query_agent_turn(self, _user, _chat_id, _turn_id): async def query_agent_timeline_events(self, _turn_id, after_sequence=0, limit=2000): return [event for event in self.persisted_events if event.sequence > after_sequence][:limit] - async def query_agent_artifacts_by_turn(self, _turn_id): - return list(self.artifacts) - class _FakeCreateTurnDbOps: def __init__(self, *, existing_turn=None, raise_on_create=False): @@ -115,16 +124,12 @@ async def create_agent_turn(self, **_kwargs): class _FakeHistoryDbOps: - def __init__(self, *, turns=None, artifacts_by_turn=None): + def __init__(self, *, turns=None): self.turns = list(turns or []) - self.artifacts_by_turn = artifacts_by_turn or {} async def query_recent_agent_turns(self, _user, _chat_id, limit=8): return self.turns[-limit:] - async def query_agent_artifacts_by_turn(self, turn_id): - return list(self.artifacts_by_turn.get(turn_id, [])) - def _build_turn(**overrides): values = { @@ -139,8 +144,6 @@ def _build_turn(**overrides): "model_profile": {"model": "gpt"}, "error_code": None, "error_message": None, - "answer_artifact_id": None, - "reference_bundle_artifact_id": None, "timeline_cursor": 1, "gmt_started": None, "gmt_finished": None, @@ -165,19 +168,6 @@ def _build_event(sequence: int, event_type: str, *, label=None, status=None, dat ) -def _build_artifact(): - return SimpleNamespace( - id="artifact-1", - turn_id="turn-1", - artifact_type="answer", - summary="answer", - payload={"text": "done"}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) - - class _FakeRequest: def __init__(self, *, base_url="http://testserver/", headers=None): self.base_url = base_url @@ -189,66 +179,32 @@ async def is_disconnected(self): @pytest.mark.asyncio async def test_turn_snapshot_returns_canonical_uimessage_parts_for_completed_turn(): - """D8.4d (#90): snapshot endpoint returns ``parts: UIMessagePart[]`` - instead of the legacy ``{turn, timeline, artifacts}`` shape. - - A completed turn with both ``answer`` and ``reference_bundle`` - artifacts should project into a ``TextPart`` followed by - ``SourceUrlPart`` / ``DataCitationPart`` pairs, with the canonical - ``status`` and ``timeline_cursor`` mirrored from the runtime - state. Per D8 §2 wire/at-rest byte-equal canonical, the FE - consumes the same discriminated union from snapshot and live SSE. + """Phase 8 D8.6 (#80) chunk-2: snapshot endpoint reads canonical + ``UIMessagePart[]`` directly from ``UIMessageStore`` (no artifact + fallback). The runtime now persists assistant text + citations as + a single ``UIMessage`` row at end-of-turn. """ - turn = _build_turn( - status=AgentTurnStatus.COMPLETED, - answer_artifact_id="artifact-1", - reference_bundle_artifact_id="bundle-1", - timeline_cursor=2, - ) - answer_artifact = SimpleNamespace( - id="artifact-1", - turn_id="turn-1", - artifact_type=AgentArtifactType.ANSWER, - summary="answer", - payload={"text": "Hello world"}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) - reference_artifact = SimpleNamespace( - id="bundle-1", - turn_id="turn-1", - artifact_type=AgentArtifactType.REFERENCE_BUNDLE, - summary="1 references", - payload={ - "items": [ - { - "source_type": "web_page", - "source_id": "src-1", - "title": "Example", - "snippet": "ApeRAG is great", - "uri": "https://example.com/a", - "metadata": {}, - } - ] - }, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), + turn = _build_turn(status=AgentTurnStatus.COMPLETED, timeline_cursor=2) + persisted = UIMessage( + id="msg-turn-1", + role="assistant", + parts=[ + TextPart(text="Hello world"), + SourceUrlPart(source_id="src-1", url="https://example.com/a", title="Example"), + DataCitationPart( + data=CitationData( + cited_text="ApeRAG is great", + location=UrlCitationLocation(url="https://example.com/a", title="Example"), + ) + ), + ], ) service = TurnService( - db_ops=_FakeDbOps( - turn=turn, - artifacts=[answer_artifact, reference_artifact], - ), - redis_store=_FakeRedisStore( - runtime_state={ - "status": "COMPLETED", - "timeline_cursor": 2, - }, - ), + db_ops=_FakeDbOps(turn=turn), + redis_store=_FakeRedisStore(runtime_state={"status": "COMPLETED", "timeline_cursor": 2}), + uimessage_store=_FakeUIMessageStore(messages={"turn-1": persisted}), ) snapshot = await service.get_turn_snapshot("user-1", "chat-1", "turn-1") @@ -270,10 +226,9 @@ async def test_turn_snapshot_returns_canonical_uimessage_parts_for_completed_tur @pytest.mark.asyncio async def test_turn_snapshot_surfaces_error_text_for_failed_turn(): - """D8.4d (#90): a FAILED turn's ``error_summary`` artifact surfaces - via the ``error_text`` field, not as a part. Mirrors the FE - transitional adapter (#77 huangheng ``snapshot-fallback.ts``) so - the BE can replace it cleanly post-merge. + """Phase 8 D8.6 (#80) chunk-2: a FAILED turn's ``error_text`` comes + straight off the ``AgentTurn`` row — the legacy ``error_summary`` + artifact was removed when ``agent_artifact`` got dropped. """ turn = _build_turn( @@ -282,26 +237,17 @@ async def test_turn_snapshot_surfaces_error_text_for_failed_turn(): error_message="upstream timeout", timeline_cursor=3, ) - error_artifact = SimpleNamespace( - id="error-1", - turn_id="turn-1", - artifact_type=AgentArtifactType.ERROR_SUMMARY, - summary="upstream timeout", - payload={"message": "Detailed: upstream timeout after 30s"}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) service = TurnService( - db_ops=_FakeDbOps(turn=turn, artifacts=[error_artifact]), + db_ops=_FakeDbOps(turn=turn), redis_store=_FakeRedisStore(runtime_state={"status": "FAILED"}), + uimessage_store=_FakeUIMessageStore(), ) snapshot = await service.get_turn_snapshot("user-1", "chat-1", "turn-1") assert snapshot.status == "FAILED" - assert snapshot.error_text == "Detailed: upstream timeout after 30s" + assert snapshot.error_text == "upstream timeout" # No assistant text was produced; error is not modelled as a part. assert snapshot.parts == [] @@ -316,6 +262,7 @@ async def test_turn_snapshot_does_not_expose_legacy_keys(): service = TurnService( db_ops=_FakeDbOps(turn=_build_turn(status=AgentTurnStatus.QUEUED)), redis_store=_FakeRedisStore(), + uimessage_store=_FakeUIMessageStore(), ) snapshot = await service.get_turn_snapshot("user-1", "chat-1", "turn-1") @@ -360,21 +307,15 @@ async def test_event_service_to_event_envelope_adds_user_activity_contract(): @pytest.mark.asyncio async def test_turn_snapshot_user_activity_inference_runs_via_event_service(): """``TurnService.get_turn_snapshot`` no longer projects timeline - events (D8.4d / #90 dropped the legacy timeline shape). The user - activity inference contract still belongs to ``EventService`` and - is exercised directly by - ``test_event_service_to_event_envelope_adds_user_activity_contract``. - - This test pins the empty-timeline guarantee: snapshots no longer - include timeline / artifacts arrays, so the FE renderer is the - sole consumer of activity hints (via the live SSE - ``data-activity`` part). + events. The user activity inference contract still belongs to + ``EventService`` and is exercised directly. """ turn = _build_turn() service = TurnService( db_ops=_FakeDbOps(turn=turn), redis_store=_FakeRedisStore(runtime_state=None), + uimessage_store=_FakeUIMessageStore(), ) snapshot = await service.get_turn_snapshot("user-1", "chat-1", "turn-1") @@ -432,30 +373,30 @@ async def _fake_get_async_client(_cls, redis_url=None): @pytest.mark.asyncio async def test_history_writer_builds_context_from_v3_turns_only(): + """Phase 8 D8.6 (#80) chunk-2: ``build_history_context`` now reads + canonical ``UIMessage`` text parts per-turn instead of legacy + ``answer`` artifact rows. Turns whose persisted message is empty + contribute nothing to the prompt context. + """ + completed_turn = _build_turn( id="turn-v3", status=AgentTurnStatus.COMPLETED, input_text="new question", - answer_artifact_id="artifact-answer", ) incomplete_answer_turn = _build_turn( id="turn-v3-missing-answer", status=AgentTurnStatus.COMPLETED, input_text="missing answer question", ) - answer_artifact = SimpleNamespace( - id="artifact-answer", - artifact_type="answer", - payload={"text": "new answer"}, + persisted_message = UIMessage( + id="msg-turn-v3", + role="assistant", + parts=[TextPart(text="new answer")], ) writer = HistoryWriter( - db_ops=_FakeHistoryDbOps( - turns=[completed_turn, incomplete_answer_turn], - artifacts_by_turn={ - "turn-v3": [answer_artifact], - "turn-v3-missing-answer": [], - }, - ) + db_ops=_FakeHistoryDbOps(turns=[completed_turn, incomplete_answer_turn]), + uimessage_store=_FakeUIMessageStore(messages={"turn-v3": persisted_message}), ) context = await writer.build_history_context("user-1", "chat-1") @@ -466,7 +407,12 @@ async def test_history_writer_builds_context_from_v3_turns_only(): @pytest.mark.asyncio -async def test_agent_runtime_views_create_stream_snapshot_cancel_and_artifact(monkeypatch): +async def test_agent_runtime_views_create_stream_snapshot_and_cancel(monkeypatch): + """Phase 8 D8.6 (#80) chunk-2: ``/agent/artifacts/{id}`` route is + deleted along with the artifact_service entry point. This test + pins the remaining create / snapshot / cancel / stream routes. + """ + turn = _build_turn(status=AgentTurnStatus.RUNNING, timeline_cursor=1) snapshot = AgentTurnSnapshot( turn_id="turn-1", @@ -492,16 +438,6 @@ async def test_agent_runtime_views_create_stream_snapshot_cancel_and_artifact(mo data={"chat_id": "chat-1"}, ) ] - artifact = AgentArtifactEnvelope( - artifact_id="artifact-1", - turn_id="turn-1", - artifact_type="answer", - summary="done", - payload={"text": "done"}, - storage_ref=None, - created_at=_now().isoformat(), - updated_at=_now().isoformat(), - ) class _FakeTurnService: def __init__(self): @@ -531,15 +467,10 @@ class _FakeEventService: async def get_events_after(self, _turn_id, after_sequence=0, limit=500): return timeline_for_stream if after_sequence == 0 else [] - class _FakeArtifactService: - async def get_artifact_for_user(self, _user, _artifact_id): - return artifact - class _FakeRuntimeManager: def __init__(self): self.turn_service = _FakeTurnService() self.event_service = _FakeEventService() - self.artifact_service = _FakeArtifactService() self.tasks = {} self.claim_turn_id = None self.claim_turn_result = "lease-owner" @@ -583,12 +514,6 @@ async def cancel_turn(self, turn_id): assert snapshot_response.turn_id == "turn-1" assert snapshot_response.role == "assistant" - artifact_response = await agent_runtime_view.get_artifact_view( - "artifact-1", - user=user, - ) - assert artifact_response.artifact_id == "artifact-1" - cancel_response = await agent_runtime_view.cancel_turn_view( "chat-1", "turn-1", @@ -618,6 +543,17 @@ async def cancel_turn(self, turn_id): assert stream_response.headers["x-vercel-ai-ui-message-stream"] == "v1" +def test_agent_runtime_views_no_artifact_route(): + """Phase 8 D8.6 (#80) chunk-2: the legacy + ``/agent/artifacts/{artifact_id}`` route + ``ArtifactService`` are + removed. Pin their absence as a regression guard. + """ + + paths = {route.path for route in agent_runtime_view.router.routes} + assert "/agent/artifacts/{artifact_id}" not in paths + assert not hasattr(agent_runtime_view, "get_artifact_view") + + @pytest.mark.asyncio async def test_agent_runtime_view_create_skips_launch_when_turn_claim_fails(monkeypatch): turn = _build_turn(status=AgentTurnStatus.QUEUED, timeline_cursor=0) diff --git a/tests/unit_test/agent_runtime/test_snapshot_assembler.py b/tests/unit_test/agent_runtime/test_snapshot_assembler.py deleted file mode 100644 index b1113b24e..000000000 --- a/tests/unit_test/agent_runtime/test_snapshot_assembler.py +++ /dev/null @@ -1,171 +0,0 @@ -# Copyright 2025 ApeCloud, Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Snapshot assembler contract tests (Phase 8 D8.4d / #90). - -Pins the legacy artifact -> ``UIMessagePart`` projection that the -snapshot endpoint relies on until the wire emitter starts persisting -``agent_message.parts`` directly (D8.6 / #80 cleanup). Mirrors the FE -side ``snapshot-fallback.ts`` adapter so deleting both sides post-D8.6 -is mechanical. -""" - -from __future__ import annotations - -from datetime import datetime, timezone -from types import SimpleNamespace - -from aperag.domains.agent_runtime.db.models import AgentArtifactType -from aperag.domains.agent_runtime.snapshot_assembler import ( - assemble_parts_from_artifacts, - extract_error_text, -) -from aperag.domains.agent_runtime.uimessage import ( - DataCitationPart, - SourceUrlPart, - TextPart, -) - - -def _now() -> datetime: - return datetime.now(timezone.utc) - - -def _answer(text: str = "Hello world") -> SimpleNamespace: - return SimpleNamespace( - id="art-answer", - turn_id="turn-1", - artifact_type=AgentArtifactType.ANSWER, - summary=text[:200], - payload={"text": text, "query": "hi"}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) - - -def _reference_bundle(items: list[dict]) -> SimpleNamespace: - return SimpleNamespace( - id="art-bundle", - turn_id="turn-1", - artifact_type=AgentArtifactType.REFERENCE_BUNDLE, - summary=f"{len(items)} references", - payload={"items": items}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) - - -def _error_summary(message: str) -> SimpleNamespace: - return SimpleNamespace( - id="art-error", - turn_id="turn-1", - artifact_type=AgentArtifactType.ERROR_SUMMARY, - summary=message[:200], - payload={"message": message}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) - - -def test_answer_artifact_projects_to_text_part(): - parts = assemble_parts_from_artifacts([_answer("ApeRAG is great")]) - assert len(parts) == 1 - assert isinstance(parts[0], TextPart) - assert parts[0].text == "ApeRAG is great" - - -def test_reference_bundle_projects_to_source_url_plus_data_citation(): - items = [ - { - "source_type": "web_page", - "source_id": "src-1", - "title": "Example", - "snippet": "ApeRAG is a RAG platform", - "uri": "https://example.com/a", - "metadata": {}, - }, - { - "source_type": "web_page", - "source_id": "src-2", - "title": "Second", - "snippet": "Second snippet", - "uri": "https://example.com/b", - "metadata": {}, - }, - ] - parts = assemble_parts_from_artifacts([_reference_bundle(items)]) - types = [getattr(part, "type", None) for part in parts] - assert types == ["source-url", "data-citation", "source-url", "data-citation"] - assert isinstance(parts[0], SourceUrlPart) - assert parts[0].source_id == "src-1" - assert parts[0].url == "https://example.com/a" - assert isinstance(parts[1], DataCitationPart) - assert parts[1].data.cited_text == "ApeRAG is a RAG platform" - assert parts[1].data.location.url == "https://example.com/a" - - -def test_reference_item_without_uri_emits_citation_only(): - """Pre-D8.6 ApeRAG reference items may be document-based with no - URL. The citation still serialises (with empty ``url``) so the FE - can show snippet/title; D8.6 will replace this with the canonical - at-rest read where citation locations carry full doc/page info. - """ - - items = [{"source_id": "doc-1", "title": "PDF chapter", "snippet": "abc"}] - parts = assemble_parts_from_artifacts([_reference_bundle(items)]) - assert [getattr(part, "type", None) for part in parts] == ["data-citation"] - assert parts[0].data.location.url == "" - assert parts[0].data.location.title == "PDF chapter" - - -def test_answer_then_reference_bundle_ordering(): - items = [{"source_id": "src-1", "uri": "https://example.com/", "title": "x", "snippet": "y"}] - parts = assemble_parts_from_artifacts([_reference_bundle(items), _answer("done")]) - types = [getattr(part, "type", None) for part in parts] - # TextPart from answer comes first regardless of input ordering. - assert types == ["text", "source-url", "data-citation"] - - -def test_unknown_artifact_types_are_skipped(): - tool_summary = SimpleNamespace( - id="art-tool", - turn_id="turn-1", - artifact_type=AgentArtifactType.TOOL_RESULT_SUMMARY, - summary="tool output", - payload={"text": "tool output"}, - storage_ref=None, - gmt_created=_now(), - gmt_updated=_now(), - ) - parts = assemble_parts_from_artifacts([tool_summary]) - assert parts == [] - - -def test_extract_error_text_prefers_payload_message(): - text = extract_error_text([_error_summary("upstream timeout")]) - assert text == "upstream timeout" - - -def test_extract_error_text_falls_back_to_summary_when_payload_empty(): - artifact = _error_summary("ignored") - artifact.payload = {} - artifact.summary = "from-summary" - assert extract_error_text([artifact]) == "from-summary" - - -def test_extract_error_text_returns_none_without_error_summary(): - assert extract_error_text([_answer()]) is None diff --git a/tests/unit_test/chat/test_chat_completion_service.py b/tests/unit_test/chat/test_chat_completion_service.py index ba7fdde20..299ec3d0d 100644 --- a/tests/unit_test/chat/test_chat_completion_service.py +++ b/tests/unit_test/chat/test_chat_completion_service.py @@ -5,6 +5,13 @@ import aperag.domains.conversation.service.chat_completion_service as completion_module from aperag.domains.agent_runtime.db.models import AgentTurnStatus +from aperag.domains.agent_runtime.uimessage import ( + CitationData, + DataCitationPart, + TextPart, + UIMessage, + UrlCitationLocation, +) from aperag.utils.constant import DOC_QA_REFERENCES @@ -33,24 +40,36 @@ def _turn(*, turn_id="turn-1", chat_id="chat-1", status=AgentTurnStatus.COMPLETE ) -def _artifacts(*, answer_text="done", references=None): - """Phase 8 D8.4d (#90): the OpenAI-compat completion path now reads - raw ``AgentArtifact`` rows directly from the DB instead of through - ``get_turn_snapshot`` (which now returns canonical UIMessage parts - for the FE).""" - references = references or [{"title": "Doc 1", "snippet": "hello"}] - return [ - SimpleNamespace( - artifact_id="artifact-answer", - artifact_type="answer", - payload={"text": answer_text}, - ), - SimpleNamespace( - artifact_id="artifact-refs", - artifact_type="reference_bundle", - payload={"items": references}, - ), - ] +def _persisted_message(*, answer_text="done", references=None) -> UIMessage: + """Phase 8 D8.6 (#80) chunk-2: the OpenAI-compat completion path + now reads canonical ``UIMessage`` parts (``TextPart`` for the + answer + ``DataCitationPart`` for each reference) directly from + ``UIMessageStore`` instead of the legacy artifact rows. + """ + + references = references or [{"title": "Doc 1", "snippet": "hello", "url": "https://example.com/doc1"}] + parts = [TextPart(text=answer_text)] + for ref in references: + parts.append( + DataCitationPart( + data=CitationData( + cited_text=ref.get("snippet", ""), + location=UrlCitationLocation( + url=ref.get("url", ""), + title=ref.get("title"), + ), + ) + ) + ) + return UIMessage(id="msg-turn-1", role="assistant", parts=parts) + + +class _FakeUIMessageStore: + def __init__(self, message: UIMessage | None): + self._message = message + + async def read(self, _turn_id): + return self._message class _FakeEventService: @@ -67,17 +86,13 @@ async def get_events_after(self, _turn_id, after_sequence=0, limit=500): class _FakeTurnService: - def __init__(self, *, chat, bot, turn, artifacts, query_turns=None): + def __init__(self, *, chat, bot, turn, query_turns=None): self.chat = chat self.bot = bot self.turn = turn - self.artifacts = list(artifacts or []) self.query_turns = list(query_turns or [turn]) self.created_requests = [] - self.db_ops = SimpleNamespace( - query_agent_turn=self._query_agent_turn, - query_agent_artifacts_by_turn=self._query_agent_artifacts_by_turn, - ) + self.db_ops = SimpleNamespace(query_agent_turn=self._query_agent_turn) async def get_chat_and_bot(self, _user, _chat_id): return self.chat, self.bot @@ -91,14 +106,12 @@ async def _query_agent_turn(self, _user, _chat_id, _turn_id): return self.query_turns.pop(0) return self.query_turns[0] - async def _query_agent_artifacts_by_turn(self, _turn_id): - return list(self.artifacts) - class _FakeRuntimeManager: - def __init__(self, *, turn_service, event_service): + def __init__(self, *, turn_service, event_service, uimessage_store): self.turn_service = turn_service self.event_service = event_service + self.uimessage_store = uimessage_store self.tasks = {} self.launch_calls = [] self.cancel_calls = [] @@ -134,9 +147,12 @@ async def test_openai_chat_completions_returns_openai_response_and_maps_override chat = SimpleNamespace(id="chat-1") bot = _bot() turn = _turn() - artifacts = _artifacts(answer_text="final answer") - fake_turn_service = _FakeTurnService(chat=chat, bot=bot, turn=turn, artifacts=artifacts) - fake_runtime_manager = _FakeRuntimeManager(turn_service=fake_turn_service, event_service=_FakeEventService()) + fake_turn_service = _FakeTurnService(chat=chat, bot=bot, turn=turn) + fake_runtime_manager = _FakeRuntimeManager( + turn_service=fake_turn_service, + event_service=_FakeEventService(), + uimessage_store=_FakeUIMessageStore(_persisted_message(answer_text="final answer")), + ) monkeypatch.setattr(completion_module, "runtime_manager", fake_runtime_manager) @@ -173,9 +189,12 @@ async def test_openai_chat_completions_creates_and_cleans_up_ephemeral_chat(monk chat = SimpleNamespace(id="chat-ephemeral") bot = _bot() turn = _turn(chat_id="chat-ephemeral") - artifacts = _artifacts(answer_text="ephemeral answer") - fake_turn_service = _FakeTurnService(chat=chat, bot=bot, turn=turn, artifacts=artifacts) - fake_runtime_manager = _FakeRuntimeManager(turn_service=fake_turn_service, event_service=_FakeEventService()) + fake_turn_service = _FakeTurnService(chat=chat, bot=bot, turn=turn) + fake_runtime_manager = _FakeRuntimeManager( + turn_service=fake_turn_service, + event_service=_FakeEventService(), + uimessage_store=_FakeUIMessageStore(_persisted_message(answer_text="ephemeral answer")), + ) fake_chat_service = _FakeChatService(created_chat_id="chat-ephemeral") monkeypatch.setattr(completion_module, "runtime_manager", fake_runtime_manager) @@ -200,18 +219,20 @@ async def test_openai_chat_completions_streams_sse_from_runtime_events(monkeypat bot = _bot() running_turn = _turn(status=AgentTurnStatus.RUNNING) completed_turn = _turn(status=AgentTurnStatus.COMPLETED) - artifacts = _artifacts(answer_text="streamed answer") fake_turn_service = _FakeTurnService( chat=chat, bot=bot, turn=running_turn, - artifacts=artifacts, query_turns=[running_turn, completed_turn], ) fake_event_service = _FakeEventService( events=[SimpleNamespace(sequence=1, type="text.delta", data={"delta": "hello"})] ) - fake_runtime_manager = _FakeRuntimeManager(turn_service=fake_turn_service, event_service=fake_event_service) + fake_runtime_manager = _FakeRuntimeManager( + turn_service=fake_turn_service, + event_service=fake_event_service, + uimessage_store=_FakeUIMessageStore(_persisted_message(answer_text="streamed answer")), + ) monkeypatch.setattr(completion_module, "runtime_manager", fake_runtime_manager) diff --git a/tests/unit_test/chat/test_chat_service.py b/tests/unit_test/chat/test_chat_service.py index e6412ad84..a69318ac0 100644 --- a/tests/unit_test/chat/test_chat_service.py +++ b/tests/unit_test/chat/test_chat_service.py @@ -19,9 +19,12 @@ from aperag.domains.agent_runtime.db.models import AgentTurnStatus from aperag.domains.agent_runtime.uimessage import ( + CitationData, DataCitationPart, SourceUrlPart, TextPart, + UIMessage, + UrlCitationLocation, ) from aperag.domains.conversation.service.chat_service import ChatService @@ -30,6 +33,14 @@ def _now(): return datetime.now(timezone.utc) +class _FakeUIMessageStore: + def __init__(self, messages=None): + self.messages = dict(messages or {}) + + async def read(self, turn_id): + return self.messages.get(turn_id) + + class _FakeChatDbOps: def __init__(self): self.chat = SimpleNamespace( @@ -47,8 +58,6 @@ def __init__(self): user="user-1", input_text="What changed?", status=AgentTurnStatus.COMPLETED, - answer_artifact_id="artifact-answer", - reference_bundle_artifact_id="artifact-refs", error_message=None, timeline_cursor=2, gmt_created=_now(), @@ -56,30 +65,6 @@ def __init__(self): gmt_finished=_now(), gmt_updated=_now(), ) - self.answer_artifact = SimpleNamespace( - id="artifact-answer", - artifact_type="answer", - summary="Here is the answer.", - payload={"text": "Here is the answer."}, - ) - self.reference_artifact = SimpleNamespace( - id="artifact-refs", - artifact_type="reference_bundle", - summary="1 references", - payload={ - "items": [ - { - "title": "Doc A", - "snippet": "Reference snippet", - "score": 0.9, - "source_type": "search_collection", - "source_id": "doc-a", - "uri": "https://example.com/doc-a", - "metadata": {"section": "intro"}, - } - ] - }, - ) async def query_chat(self, user, bot_id, chat_id): if user == "user-1" and bot_id == "bot-1" and chat_id == "chat-1": @@ -91,23 +76,32 @@ async def query_agent_turns(self, user, chat_id): return [self.turn] return [] - async def query_agent_artifacts_by_turn(self, turn_id): - if turn_id == "turn-1": - return [self.answer_artifact, self.reference_artifact] - return [] + +def _build_persisted_message() -> UIMessage: + return UIMessage( + id="msg-turn-1", + role="assistant", + parts=[ + TextPart(text="Here is the answer."), + SourceUrlPart(source_id="doc-a", url="https://example.com/doc-a", title="Doc A"), + DataCitationPart( + data=CitationData( + cited_text="Reference snippet", + location=UrlCitationLocation(url="https://example.com/doc-a", title="Doc A"), + ) + ), + ], + ) @pytest.mark.asyncio async def test_get_chat_returns_canonical_uimessage_history(): - """Phase 8 D8.5-BE (#92): ``ChatDetails.history`` is now a list of - canonical ``AgentTurnSnapshot`` envelopes (one per assistant turn), - each carrying the same ``UIMessagePart`` shape the FE consumes from - the live SSE stream. Replaces the legacy - ``list[list[ChatMessage]]`` shape so historical and live turns - render through a single canonical path. + """Phase 8 D8.6 (#80) chunk-2: ``ChatDetails.history`` reads + canonical ``UIMessage`` parts directly from ``UIMessageStore`` — + legacy ``AgentArtifact`` projection is gone. """ - service = ChatService() + service = ChatService(uimessage_store=_FakeUIMessageStore({"turn-1": _build_persisted_message()})) service.db_ops = _FakeChatDbOps() chat = await service.get_chat("user-1", "bot-1", "chat-1") @@ -143,11 +137,13 @@ async def test_get_chat_returns_canonical_uimessage_history(): @pytest.mark.asyncio async def test_get_chat_history_surfaces_error_text_for_failed_turn(): - """A FAILED turn's error_text comes from ``error_summary`` artifact - when present, falling back to ``turn.error_message`` (mirrors the - snapshot endpoint contract from #90 D8.4d).""" + """Phase 8 D8.6 (#80) chunk-2: a FAILED turn's ``error_text`` comes + straight off the ``AgentTurn`` row (the legacy ``error_summary`` + artifact is gone). ``parts`` is empty when the runtime never + persisted a UIMessage for the failed turn. + """ - service = ChatService() + service = ChatService(uimessage_store=_FakeUIMessageStore()) db_ops = _FakeChatDbOps() db_ops.turn = SimpleNamespace( id="turn-failed", @@ -155,8 +151,6 @@ async def test_get_chat_history_surfaces_error_text_for_failed_turn(): user="user-1", input_text="Trigger failure", status=AgentTurnStatus.FAILED, - answer_artifact_id=None, - reference_bundle_artifact_id=None, error_message="upstream provider timeout", timeline_cursor=1, gmt_created=_now(), @@ -164,13 +158,6 @@ async def test_get_chat_history_surfaces_error_text_for_failed_turn(): gmt_finished=_now(), gmt_updated=_now(), ) - db_ops.answer_artifact = None - db_ops.reference_artifact = None - - async def _empty_artifacts(turn_id): - return [] - - db_ops.query_agent_artifacts_by_turn = _empty_artifacts service.db_ops = db_ops chat = await service.get_chat("user-1", "bot-1", "chat-1") @@ -192,7 +179,7 @@ async def test_get_chat_history_does_not_expose_legacy_chatmessage_shape(): the canonical snapshot shape. """ - service = ChatService() + service = ChatService(uimessage_store=_FakeUIMessageStore({"turn-1": _build_persisted_message()})) service.db_ops = _FakeChatDbOps() chat = await service.get_chat("user-1", "bot-1", "chat-1")