add option to stream holmes chat function

Author: RoiGlinik

src/robusta/core/model/base_params.py

@@ -2,7 +2,7 @@
 from enum import Enum, auto
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, SecretStr, validator
+from pydantic import BaseModel, SecretStr, validator, Field
 
 from robusta.integrations import openshift
 from robusta.utils.documented_pydantic import DocumentedModel
@@ -190,6 +190,7 @@ class HolmesChatParams(HolmesParams):
     ask: str
     conversation_history: Optional[list[dict]] = None
     render_graph_images: bool = False
+    stream: bool = Field(default=False)
 
 
 class HolmesIssueChatParams(HolmesChatParams):
@@ -251,7 +252,6 @@ class HolmesWorkloadHealthChatParams(HolmesParams):
     conversation_history: Optional[list[dict]] = None
 
 
-
 class NamespacedResourcesParams(ActionParams):
     """
     :var name: Resource name

src/robusta/core/playbooks/internal/ai_integration.py

@@ -351,8 +351,24 @@ def holmes_chat(event: ExecutionBaseEvent, params: HolmesChatParams):
     cluster_name = event.get_context().cluster_name
 
     try:
-        holmes_req = HolmesChatRequest(ask=params.ask, conversation_history=params.conversation_history, model=params.model)
-        result = requests.post(f"{holmes_url}/api/chat", data=holmes_req.json())
+        holmes_req = HolmesChatRequest(ask=params.ask, conversation_history=params.conversation_history, model=params.model, stream=params.stream)
+        url = f"{holmes_url}/api/chat"
+        if params.stream:
+            with requests.post(
+                url,
+                data=holmes_req.json(),
+                stream=True,
+                headers={"Connection": "keep-alive"},
+            ) as resp:
+                resp.raise_for_status()
+                for line in resp.iter_content(
+                    chunk_size=None, decode_unicode=True
+                ):  # Avoid streaming chunks from holmes. send them as they arrive.
+                    if line:
+                        event.ws(data=line)
+            return
+
+        result = requests.post(url, data=holmes_req.json())
         result.raise_for_status()
         holmes_result = HolmesChatResult(**json.loads(result.text))
         holmes_result.files = []

src/robusta/core/reporting/holmes.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from robusta.core.model.base_params import (
     ConversationType,
@@ -38,6 +38,7 @@ class HolmesChatRequest(BaseModel):
     ask: str
     conversation_history: Optional[List[dict]] = None
     model: Optional[str] = None
+    stream: bool = Field(default=False)
 
 
 class HolmesIssueChatRequest(HolmesChatRequest):