fix comments

Author: guglielmo-san

src/a2a/helpers/proto_helpers.py

@@ -64,10 +64,11 @@ def new_artifact(
     parts: list[Part],
     name: str,
     description: str | None = None,
+    artifact_id: str | None = None,
 ) -> Artifact:
     """Creates a new Artifact object."""
     return Artifact(
-        artifact_id=str(uuid.uuid4()),
+        artifact_id=artifact_id or str(uuid.uuid4()),
         parts=parts,
         name=name,
         description=description,
@@ -78,12 +79,14 @@ def new_text_artifact(
     name: str,
     text: str,
     description: str | None = None,
+    artifact_id: str | None = None,
 ) -> Artifact:
     """Creates a new Artifact object containing only a single text Part."""
     return new_artifact(
         [Part(text=text)],
         name,
         description,
+        artifact_id=artifact_id,
     )
 
 
@@ -97,6 +100,8 @@ def get_artifact_text(artifact: Artifact, delimiter: str = '\n') -> str:
 
 def new_task_from_user_message(user_message: Message) -> Task:
     """Creates a new Task object from an initial user message."""
+    if user_message.role != Role.ROLE_USER:
+        raise ValueError('Message must be from a user')
     if not user_message.parts:
         raise ValueError('Message parts cannot be empty')
     for part in user_message.parts:
@@ -173,13 +178,14 @@ def new_text_artifact_update_event(  # noqa: PLR0913
     text: str,
     append: bool = False,
     last_chunk: bool = False,
+    artifact_id: str | None = None,
 ) -> TaskArtifactUpdateEvent:
     """Creates a TaskArtifactUpdateEvent with a single text artifact."""
     return TaskArtifactUpdateEvent(
         task_id=task_id,
         context_id=context_id,
-        artifact=Artifact(
-            artifact_id=str(uuid.uuid4()), name=name, parts=[Part(text=text)]
+        artifact=new_text_artifact(
+            name=name, text=text, artifact_id=artifact_id
         ),
         append=append,
         last_chunk=last_chunk,

tests/helpers/test_proto_helpers.py

@@ -79,6 +79,15 @@ def test_new_text_artifact() -> None:
     assert art.artifact_id != ''
 
 
+def test_new_text_artifact_with_id() -> None:
+    art = new_text_artifact(name='test', text='content', description='desc', artifact_id='art1')
+    assert art.name == 'test'
+    assert art.description == 'desc'
+    assert len(art.parts) == 1
+    assert art.parts[0].text == 'content'
+    assert art.artifact_id == 'art1'
+
+
 def test_get_artifact_text() -> None:
     art = Artifact(parts=[Part(text='hello'), Part(text='world')])
     assert get_artifact_text(art) == 'hello\nworld'
@@ -171,6 +180,21 @@ def test_new_text_artifact_update_event() -> None:
     assert event.last_chunk is True
 
 
+def test_new_text_artifact_update_event_with_id() -> None:
+    event = new_text_artifact_update_event(
+        task_id='task1',
+        context_id='ctx1',
+        name='test',
+        text='content',
+        artifact_id='art1',
+    )
+    assert event.task_id == 'task1'
+    assert event.context_id == 'ctx1'
+    assert event.artifact.name == 'test'
+    assert event.artifact.parts[0].text == 'content'
+    assert event.artifact.artifact_id == 'art1'
+
+
 def test_get_stream_response_text_message() -> None:
     resp = StreamResponse(message=Message(parts=[Part(text='hello')]))
     assert get_stream_response_text(resp) == 'hello'