Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 39 additions & 1 deletion gooddata-sdk/gooddata_sdk/compute/service.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
# (C) 2022 GoodData Corporation
from __future__ import annotations

import json
import logging
from collections.abc import Iterator
from typing import Any, Optional

from gooddata_api_client import ApiException
Expand Down Expand Up @@ -89,12 +91,48 @@ def ai_chat(self, workspace_id: str, question: str) -> ChatResult:
workspace_id: workspace identifier
question: question to ask AI
Returns:
str: Chat response
ChatResult: Chat response
"""
chat_request = ChatRequest(question=question)
response = self._actions_api.ai_chat(workspace_id, chat_request, _check_return_type=False)
return response

def _parse_sse_events(self, raw: str) -> Iterator[Any]:
"""Helper to parse SSE events and yield JSON from data lines."""
events = raw.split("\n\n")
for event in events:
for line in event.split("\n"):
if line.startswith("data:"):
try:
yield json.loads(line[5:].strip())
except json.JSONDecodeError:
continue

def ai_chat_stream(self, workspace_id: str, question: str) -> Iterator[Any]:
"""
Chat Stream with AI in GoodData workspace.

Args:
workspace_id: workspace identifier
question: question to ask AI
Returns:
Iterator[Any]: Yields parsed JSON objects from each SSE event's data field
"""
chat_request = ChatRequest(question=question)
response = self._actions_api.ai_chat_stream(
workspace_id, chat_request, _check_return_type=False, _preload_content=False
)
buffer = ""
try:
for chunk in response.stream(decode_content=True):
if chunk:
buffer += chunk.decode("utf-8")
*events, buffer = buffer.split("\n\n")
for event in events:
yield from self._parse_sse_events(event)
finally:
response.release_conn()

def get_ai_chat_history(
self,
workspace_id: str,
Expand Down
Loading
Loading