Add prompt suggestions SSE helper

README.md

@@ -9,6 +9,7 @@ For some example agents that demonstrate the full usage of the SDK, see the
 ## Features
 
 - [Streaming Conversations](#message_chunk)
+- [Follow-up Prompt Suggestions](#prompt_suggestions)
 - [Reasoning steps / status updates](#reasoning_step)
 - [Retrieve widget data from OpenBB Workspace](#get_widget_data)
 - [Citations](#cite-and-citations)
@@ -136,6 +137,18 @@ yield reasoning_step(
 ).model_dump()
 ```
 
+### `prompt_suggestions`
+
+Send follow-up prompt suggestions to OpenBB Workspace after an agent response.
+
+```python
+from openbb_ai.helpers import prompt_suggestions
+
+yield prompt_suggestions(
+    ["Summarize the main takeaways", "Compare revenue and margin trends"]
+).model_dump()
+```
+
 ### `get_widget_data`
 
 Create a function call SSE that retrieves data from widgets on the OpenBB

openbb_ai/__init__.py

@@ -3,6 +3,7 @@
 from .helpers import cite as cite
 from .helpers import get_widget_data as get_widget_data
 from .helpers import message_chunk as message_chunk
+from .helpers import prompt_suggestions as prompt_suggestions
 from .helpers import reasoning_step as reasoning_step
 from .helpers import table as table
 from .models import QueryRequest as QueryRequest

openbb_ai/helpers.py

@@ -17,6 +17,8 @@
     MessageChunkSSE,
     MessageChunkSSEData,
     PieChartParameters,
+    PromptSuggestionsSSE,
+    PromptSuggestionsSSEData,
     ScatterChartParameters,
     SourceInfo,
     StatusUpdateSSE,
@@ -84,6 +86,25 @@ def message_chunk(text: str) -> MessageChunkSSE:
     return MessageChunkSSE(data=MessageChunkSSEData(delta=text))
 
 
+def prompt_suggestions(suggestions: list[str]) -> PromptSuggestionsSSE:
+    """Create a prompt suggestions SSE.
+
+    This SSE is used to send follow-up prompt suggestions to OpenBB Workspace
+    after an agent response.
+
+    Parameters
+    ----------
+    suggestions: list[str]
+        Prompt suggestions to display to the user.
+
+    Returns
+    -------
+    PromptSuggestionsSSE
+        The prompt suggestions SSE.
+    """
+    return PromptSuggestionsSSE(data=PromptSuggestionsSSEData(suggestions=suggestions))
+
+
 def get_widget_data(widget_requests: list[WidgetRequest]) -> FunctionCallSSE:
     """Create a function call that retrieve data for a widget on the OpenBB Workspace
 

openbb_ai/models.py

@@ -673,6 +673,23 @@ class DashboardInfo(BaseModel):
     )
 
 
+class AgentFeatureSelectOption(BaseModel):
+    label: str
+    value: str
+
+
+class AgentFeatureOption(BaseModel):
+    label: str
+    type: Literal["toggle", "text", "select"] | None = None
+    default: bool | str | None = None
+    description: str | None = None
+    placeholder: str | None = None
+    options: list[AgentFeatureSelectOption] | None = None
+
+
+AgentFeature = bool | AgentFeatureOption
+
+
 class WorkspaceAgent(BaseModel):
     holder_url: str | None = Field(
         default=None,
@@ -689,7 +706,7 @@ class WorkspaceAgent(BaseModel):
     description: str | None = Field(
         default=None, description="A description of the agent."
     )
-    features: dict[str, bool] = Field(
+    features: dict[str, AgentFeature] = Field(
         default_factory=dict,
         description="A dictionary of features that the agent supports.",
     )
@@ -865,6 +882,15 @@ class MessageChunkSSE(BaseSSE):
     data: MessageChunkSSEData
 
 
+class PromptSuggestionsSSEData(BaseModel):
+    suggestions: list[str]
+
+
+class PromptSuggestionsSSE(BaseSSE):
+    event: Literal["copilotPromptSuggestions"] = "copilotPromptSuggestions"
+    data: PromptSuggestionsSSEData
+
+
 class MessageArtifactSSE(BaseSSE):
     event: Literal["copilotMessageArtifact"] = "copilotMessageArtifact"
     data: ClientArtifact
@@ -917,15 +943,15 @@ class StatusUpdateSSEData(BaseModel):
     @model_validator(mode="before")
     @classmethod
     def exclude_fields(cls, values):
-        # Exclude these fields from being in the "details" field.  (since this
-        # pollutes the JSON output)
+        # Exclude these fields from being in the "details" field.
+        # (since this pollutes the JSON output)
         _exclude_fields = EXCLUDE_STATUS_UPDATE_DETAILS_FIELDS
         if details := values.get("details"):
             if isinstance(details, list):
                 for detail in details:
                     if isinstance(detail, dict):
                         for key in list(detail.keys()):
-                            if key.lower() in _exclude_fields:
+                            if str(key).lower() in _exclude_fields:
                                 detail.pop(key, None)
         return values
 
@@ -937,6 +963,7 @@ class StatusUpdateSSE(BaseSSE):
 
 SSE = (
     MessageChunkSSE
+    | PromptSuggestionsSSE
     | MessageArtifactSSE
     | FunctionCallSSE
     | StatusUpdateSSE

openbb_ai/testing.py

@@ -56,6 +56,13 @@ def parse_event_stream(self):
                 self.events.append(
                     CopilotEvent(event_type=event_name, content=data_dict_)
                 )
+            elif event_type == "copilotPromptSuggestions" and line.startswith("data:"):
+                event_name = "copilotPromptSuggestions"
+                data_payload = line.split("data:")[1].strip()
+                data_dict_ = json.loads(data_payload)
+                self.events.append(
+                    CopilotEvent(event_type=event_name, content=data_dict_)
+                )
             elif event_type == "copilotCitationCollection" and line.startswith("data:"):
                 event_name = "copilotCitationCollection"
                 data_payload = line.split("data:")[1].strip()
@@ -86,6 +93,14 @@ def citations(self) -> list[CopilotEvent]:
             if event.event_type == "copilotCitationCollection"
         ]
 
+    @property
+    def prompt_suggestions(self) -> list[CopilotEvent]:
+        return [
+            event
+            for event in self.events
+            if event.event_type == "copilotPromptSuggestions"
+        ]
+
     def __iter__(self):
         return self
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "openbb-ai"
-version = "2.0.2"
+version = "2.1.2"
 description = "An SDK for building agents compatible with OpenBB Workspace"
 authors = [
     { name = "OpenBB Team", email = "hello@openbb.finance" },

tests/test_helpers.py

@@ -1,5 +1,10 @@
-from openbb_ai.helpers import chart, reasoning_step, table
-from openbb_ai.models import ClientArtifact, MessageArtifactSSE, StatusUpdateSSE
+from openbb_ai.helpers import chart, prompt_suggestions, reasoning_step, table
+from openbb_ai.models import (
+    ClientArtifact,
+    MessageArtifactSSE,
+    PromptSuggestionsSSE,
+    StatusUpdateSSE,
+)
 
 
 def test_reasoning_step():
@@ -19,6 +24,14 @@ def test_reasoning_step():
     assert result.data.details == []
 
 
+def test_prompt_suggestions():
+    result = prompt_suggestions(["Show revenue growth", "Compare margin trends"])
+
+    assert isinstance(result, PromptSuggestionsSSE)
+    assert result.event == "copilotPromptSuggestions"
+    assert result.data.suggestions == ["Show revenue growth", "Compare margin trends"]
+
+
 def test_chart_line():
     result = chart(
         type="line",

tests/test_models.py

@@ -1,6 +1,32 @@
 import uuid
 
-from openbb_ai.models import Citation, CitationHighlightBoundingBox, SourceInfo
+from openbb_ai.models import (
+    AgentFeatureOption,
+    Citation,
+    CitationHighlightBoundingBox,
+    SourceInfo,
+    WorkspaceAgent,
+)
+
+
+def test_workspace_agent_supports_feature_option_metadata():
+    agent = WorkspaceAgent(
+        id="openbb_ada",
+        name="OpenBB Copilot",
+        features={
+            "streaming": True,
+            "prompt-suggestions": {
+                "label": "Follow-up Suggestions",
+                "default": True,
+                "description": "Show follow-up prompt suggestions after each response.",
+            },
+        },
+    )
+
+    feature = agent.features["prompt-suggestions"]
+    assert isinstance(feature, AgentFeatureOption)
+    assert feature.label == "Follow-up Suggestions"
+    assert feature.default is True
 
 
 def test_citation_eq():