multiturnresponse for injecting assistant prefill

Author: kallsyms

asimov/graph/__init__.py

@@ -708,6 +708,8 @@ async def run_node(self, node_name: str) -> Dict[str, Any]:
                     continue
 
                 return result
+            except asyncio.CancelledError:
+                raise
             except Exception as e:
                 last_exception = e
 
@@ -922,5 +924,6 @@ async def run_from_snapshot(self, snapshot_dir: pathlib.Path):
             task = pickle.load(f)
         await self._run_task(task)
 
+
 # Need to build after CompositeModule is defined
 AgentModule.model_rebuild()

asimov/services/inference_clients.py

@@ -64,6 +64,15 @@ class ChatMessage(AsimovBase):
     model_families: List[str] = Field(default_factory=list)
 
 
+class MultiTurnResponse:
+    def __init__(self, first_response: str, others: list[ChatMessage]):
+        self.first_response = first_response
+        self.others = others
+
+    def __str__(self):
+        return self.first_response
+
+
 class AnthropicRequest(AsimovBase):
     anthropic_version: str
     system: str
@@ -342,6 +351,20 @@ async def tool_chain(
                     "content": content_blocks,
                 }
             )
+            if isinstance(result, MultiTurnResponse):
+                for other in result.others:
+                    serialized_messages.append(
+                        {
+                            "role": other.role.value,
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": other.content,
+                                }
+                            ],
+                        }
+                    )
+
         return serialized_messages