diff --git a/README.md b/README.md index 6e23143..9bb01f5 100644 --- a/README.md +++ b/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 diff --git a/openbb_ai/__init__.py b/openbb_ai/__init__.py index bcc92b8..ad06923 100644 --- a/openbb_ai/__init__.py +++ b/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 diff --git a/openbb_ai/helpers.py b/openbb_ai/helpers.py index f1de560..0cf7ee6 100644 --- a/openbb_ai/helpers.py +++ b/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 diff --git a/openbb_ai/models.py b/openbb_ai/models.py index d5705d3..e365d32 100644 --- a/openbb_ai/models.py +++ b/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 diff --git a/openbb_ai/testing.py b/openbb_ai/testing.py index 059a12b..65ec265 100644 --- a/openbb_ai/testing.py +++ b/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 diff --git a/poetry.lock b/poetry.lock index f63449d..c83f57a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -39,14 +39,14 @@ files = [ [[package]] name = "distlib" -version = "0.4.0" +version = "0.4.2" description = "Distribution utilities" optional = false python-versions = "*" groups = ["dev"] files = [ - {file = "distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16"}, - {file = "distlib-0.4.0.tar.gz", hash = "sha256:feec40075be03a04501a973d81f633735b4b69f98b05450592310c0f401a4e0d"}, + {file = "distlib-0.4.2-py2.py3-none-any.whl", hash = "sha256:ca4cb11e5d746b5ec13c199cbf19ae27a241f89702b54e153a74332955446067"}, + {file = "distlib-0.4.2.tar.gz", hash = "sha256:baeb401c90f27acd15c4861ae0847d1e731c27ac3dbf4210643ba61fa1e813db"}, ] [[package]] @@ -70,14 +70,14 @@ test = ["pytest (>=6)"] [[package]] name = "filelock" -version = "3.29.0" +version = "3.29.1" description = "A platform independent file lock." optional = false python-versions = ">=3.10" groups = ["dev"] files = [ - {file = "filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258"}, - {file = "filelock-3.29.0.tar.gz", hash = "sha256:69974355e960702e789734cb4871f884ea6fe50bd8404051a3530bc07809cf90"}, + {file = "filelock-3.29.1-py3-none-any.whl", hash = "sha256:85199dfd706869641b72b2e8955d5416a4b2b7dc4b0e8e6d97b4cc1299a6983b"}, + {file = "filelock-3.29.1.tar.gz", hash = "sha256:d97e6b1b9757569626c58caa07dc4beb1613f4a2938b1e8cc81afca398906c9e"}, ] [[package]] @@ -335,14 +335,14 @@ re2 = ["google-re2 (>=1.1)"] [[package]] name = "platformdirs" -version = "4.9.6" +version = "4.10.0" description = "A small Python package for determining appropriate platform-specific dirs, e.g. a `user data dir`." optional = false python-versions = ">=3.10" groups = ["dev"] files = [ - {file = "platformdirs-4.9.6-py3-none-any.whl", hash = "sha256:e61adb1d5e5cb3441b4b7710bea7e4c12250ca49439228cc1021c00dcfac0917"}, - {file = "platformdirs-4.9.6.tar.gz", hash = "sha256:3bfa75b0ad0db84096ae777218481852c0ebc6c727b3168c1b9e0118e458cf0a"}, + {file = "platformdirs-4.10.0-py3-none-any.whl", hash = "sha256:fb516cdb12eb0d857d0cd85a7c57cea4d060bee4578d6cf5a14dfdf8cbf8784a"}, + {file = "platformdirs-4.10.0.tar.gz", hash = "sha256:31e761a6a0ca04faf7353ea759bdba55652be214725111e5aac52dfa29d4bef7"}, ] [[package]] @@ -576,14 +576,14 @@ dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "requests [[package]] name = "python-discovery" -version = "1.3.1" +version = "1.4.0" description = "Python interpreter discovery" optional = false python-versions = ">=3.8" groups = ["dev"] files = [ - {file = "python_discovery-1.3.1-py3-none-any.whl", hash = "sha256:ed188687ebb3b82c01a17cd5ac62fc94d9f6487a7f1a0f9dfe89753fec91039c"}, - {file = "python_discovery-1.3.1.tar.gz", hash = "sha256:62f6db28064c9613e7ca76cb3f00c38c839a07c31c00dfe7ed0986493d2150a6"}, + {file = "python_discovery-1.4.0-py3-none-any.whl", hash = "sha256:26ed78d703e234879a66244c7d4114563fb13ec5cd30a2d1357e5fb4850782da"}, + {file = "python_discovery-1.4.0.tar.gz", hash = "sha256:eb8bc7daad3c226c147e45bb4e970a1feb1bf4048ee178e6db59e197b8010ce3"}, ] [package.dependencies] @@ -792,21 +792,21 @@ typing-extensions = ">=4.12.0" [[package]] name = "virtualenv" -version = "21.3.3" +version = "21.4.2" description = "Virtual Python Environment builder" optional = false python-versions = ">=3.8" groups = ["dev"] files = [ - {file = "virtualenv-21.3.3-py3-none-any.whl", hash = "sha256:7d5987d8369e098e41406efb780a3d4ca79280097293899e351a6407ee153ab3"}, - {file = "virtualenv-21.3.3.tar.gz", hash = "sha256:f5bda277e553b1c2b3c1a8debfc30496e1288cc93ce6b7b71b3280047e317328"}, + {file = "virtualenv-21.4.2-py3-none-any.whl", hash = "sha256:854210ca524a1a4d0d744734f4acbc721c3ffe163b85bbf5d56d14d5ae2f0fae"}, + {file = "virtualenv-21.4.2.tar.gz", hash = "sha256:38e6ee0a555615c0ea9da2ac7e9998fe8dc3b911dd33ad8eaad2020957653b0c"}, ] [package.dependencies] distlib = ">=0.3.7,<1" filelock = {version = ">=3.24.2,<4", markers = "python_version >= \"3.10\""} platformdirs = ">=3.9.1,<5" -python-discovery = ">=1.3.1" +python-discovery = ">=1.4" typing-extensions = {version = ">=4.13.2", markers = "python_version < \"3.11\""} [[package]] diff --git a/pyproject.toml b/pyproject.toml index 7c908eb..3f22b1e 100644 --- a/pyproject.toml +++ b/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" }, diff --git a/tests/test_helpers.py b/tests/test_helpers.py index c9e62cf..cf190f0 100644 --- a/tests/test_helpers.py +++ b/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", diff --git a/tests/test_models.py b/tests/test_models.py index 3f3d336..ee072cb 100644 --- a/tests/test_models.py +++ b/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():