add transform files

litellm/llms/a2a/chat/streaming_iterator.py

@@ -0,0 +1,103 @@
+"""
+A2A Streaming Response Iterator
+"""
+from typing import Optional, Union
+
+from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
+
+from ..common_utils import extract_text_from_a2a_response
+
+
+class A2AModelResponseIterator(BaseModelResponseIterator):
+    """
+    Iterator for parsing A2A streaming responses.
+    
+    Converts A2A JSON-RPC streaming chunks to OpenAI-compatible format.
+    """
+    
+    def __init__(
+        self,
+        streaming_response,
+        sync_stream: bool,
+        json_mode: Optional[bool] = False,
+        model: str = "a2a/agent",
+    ):
+        super().__init__(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
+        self.model = model
+    
+    def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]:
+        """
+        Parse A2A streaming chunk to OpenAI format.
+        
+        A2A chunk format:
+        {
+            "jsonrpc": "2.0",
+            "id": "request-id",
+            "result": {
+                "message": {
+                    "parts": [{"kind": "text", "text": "content"}]
+                }
+            }
+        }
+        
+        Or for tasks:
+        {
+            "jsonrpc": "2.0",
+            "result": {
+                "kind": "task",
+                "status": {"state": "running"},
+                "artifacts": [{"parts": [{"kind": "text", "text": "content"}]}]
+            }
+        }
+        """
+        try:
+            # Extract text from A2A response
+            text = extract_text_from_a2a_response(chunk)
+            
+            # Determine finish reason
+            finish_reason = self._get_finish_reason(chunk)
+            
+            # Return generic streaming chunk
+            return GenericStreamingChunk(
+                text=text,
+                is_finished=bool(finish_reason),
+                finish_reason=finish_reason or "",
+                usage=None,
+                index=0,
+                tool_use=None,
+            )
+        except Exception:
+            # Return empty chunk on parse error
+            return GenericStreamingChunk(
+                text="",
+                is_finished=False,
+                finish_reason="",
+                usage=None,
+                index=0,
+                tool_use=None,
+            )
+    
+    def _get_finish_reason(self, chunk: dict) -> Optional[str]:
+        """Extract finish reason from A2A chunk"""
+        result = chunk.get("result", {})
+        
+        # Check for task completion
+        if isinstance(result, dict):
+            status = result.get("status", {})
+            if isinstance(status, dict):
+                state = status.get("state")
+                if state == "completed":
+                    return "stop"
+                elif state == "failed":
+                    return "error"
+        
+        # Check for [DONE] marker
+        if chunk.get("done") is True:
+            return "stop"
+        
+        return None