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
64 changes: 64 additions & 0 deletions 37-vanilla-agent-custom-features/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
# Vanilla Agent with Custom Features

This is a simple example agent that demonstrates how to detect and report which custom features are enabled or disabled in the OpenBB Workspace UI.

## What it does

This agent:
- Greets users with a friendly hello message
- Reports the status of custom features (Deep Research and Web Search)
- Shows whether each feature is enabled (✅) or disabled (❌) based on the UI settings

## Features

The agent defines two custom features in its configuration:
- **Deep Research**: Allows the agent to perform deep research (default: disabled)
- **Web Search**: Allows the agent to search the web (default: enabled)

Users can toggle these features on/off in the OpenBB Workspace UI, and the agent will detect and report the current status.

## Getting started

### Prerequisites

Ensure you have poetry, a tool for dependency management and packaging in Python, as well as your OpenAI API key.

### Installation and Running

1. Clone this repository to your local machine.

2. Set the OpenAI API key as an environment variable in your .bashrc or .zshrc file:

``` sh
# in .zshrc or .bashrc
export OPENAI_API_KEY=<your-api-key>
```

3. Install the necessary dependencies:

``` sh
poetry install --no-root
```

4. Start the API server:

``` sh
cd 37-vanilla-agent-custom-features
poetry run uvicorn vanilla_agent_custom_features.main:app --port 7777 --reload
```

This command runs the FastAPI application, making it accessible on your network.

### Testing the Agent

The example agent has a small, basic test suite to ensure it's working correctly. As you develop your agent, you are highly encouraged to expand these tests.

You can run the tests with:

```sh
pytest tests
```

### Accessing the Documentation

Once the API server is running, you can view the documentation and interact with the API by visiting: http://localhost:7777/docs
Empty file.
153 changes: 153 additions & 0 deletions 37-vanilla-agent-custom-features/tests/test_agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
import json
from unittest import mock
from fastapi.testclient import TestClient
from vanilla_agent_custom_features.main import app
import pytest
from pathlib import Path
from openbb_ai.testing import CopilotResponse

test_client = TestClient(app)


@pytest.fixture(autouse=True)
def reset_sse_starlette_appstatus_event():
"""
Fixture that resets the appstatus event in the sse_starlette app.
Should be used on any test that uses sse_starlette to stream events.
"""
# See https://github.com/sysid/sse-starlette/issues/59
from sse_starlette.sse import AppStatus

AppStatus.should_exit_event = None


def test_agents_json():
"""Test that the agents.json endpoint returns the correct configuration."""
response = test_client.get("/agents.json")
assert response.status_code == 200

data = response.json()
assert "vanilla_agent_custom_features" in data

agent_config = data["vanilla_agent_custom_features"]
assert agent_config["name"] == "Vanilla Agent Custom Features"
assert "deep-research" in agent_config["features"]
assert "web-search" in agent_config["features"]
assert not agent_config["features"]["deep-research"]["default"]
assert agent_config["features"]["web-search"]["default"]


def test_query_simple_message():
"""Test basic query with a simple message."""
test_payload_path = (
Path(__file__).parent.parent.parent
/ "testing"
/ "test_payloads"
/ "single_message.json"
)
test_payload = json.load(open(test_payload_path))

response = test_client.post("/v1/query", json=test_payload)
assert response.status_code == 200
copilot_response = CopilotResponse(response.text)
assert copilot_response.has_any("copilotMessage", "2")


def test_query_with_workspace_options():
"""Test query with workspace options to verify feature detection."""
test_payload = {
"messages": [{"role": "human", "content": "Hello, what features are enabled?"}],
"workspace_options": {"deep-research": True, "web-search": False},
}

with mock.patch("openai.AsyncOpenAI") as mock_openai:
# Mock the OpenAI client
mock_client = mock.MagicMock()
mock_openai.return_value = mock_client

# Verify that the system message contains the correct feature status
async def mock_create(**kwargs):
messages = kwargs["messages"]
system_message = messages[0]["content"]

# Check that the system message contains the feature status
assert "Deep Research: ✅ Enabled" in system_message
assert "Web Search: ❌ Disabled" in system_message

# Return a mock stream
class MockEvent:
class Choice:
class Delta:
content = "Hello! Features are configured."

delta = Delta()

choices = [Choice()]

yield MockEvent()

mock_client.chat.completions.create = mock_create

response = test_client.post("/v1/query", json=test_payload)
assert response.status_code == 200


def test_query_default_workspace_options():
"""Test query without workspace options uses defaults."""
test_payload = {"messages": [{"role": "human", "content": "Hi there!"}]}

with mock.patch("openai.AsyncOpenAI") as mock_openai:
# Mock the OpenAI client
mock_client = mock.MagicMock()
mock_openai.return_value = mock_client

# Verify that the system message contains the default feature status
async def mock_create(**kwargs):
messages = kwargs["messages"]
system_message = messages[0]["content"]

# Check that defaults are used (deep-research: False, web-search: True)
assert "Deep Research: ❌ Disabled" in system_message
assert "Web Search: ✅ Enabled" in system_message

# Return a mock stream
class MockEvent:
class Choice:
class Delta:
content = "Hello!"

delta = Delta()

choices = [Choice()]

yield MockEvent()

mock_client.chat.completions.create = mock_create

response = test_client.post("/v1/query", json=test_payload)
assert response.status_code == 200


def test_query_conversation():
"""Test query with multiple messages in conversation."""
test_payload_path = (
Path(__file__).parent.parent.parent
/ "testing"
/ "test_payloads"
/ "multiple_messages.json"
)
test_payload = json.load(open(test_payload_path))

response = test_client.post("/v1/query", json=test_payload)
assert response.status_code == 200
copilot_response = CopilotResponse(response.text)
assert copilot_response.has_any("copilotMessage", "4")


def test_query_no_messages():
"""Test query with empty messages list."""
test_payload = {
"messages": [],
}
response = test_client.post("/v1/query", json=test_payload)
assert "messages list cannot be empty" in response.text
118 changes: 118 additions & 0 deletions 37-vanilla-agent-custom-features/vanilla_agent_custom_features/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
from typing import AsyncGenerator

import openai
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from openai.types.chat import (
ChatCompletionAssistantMessageParam,
ChatCompletionMessageParam,
ChatCompletionSystemMessageParam,
ChatCompletionUserMessageParam,
)
from openbb_ai import message_chunk
from openbb_ai.models import MessageChunkSSE, QueryRequest
from sse_starlette.sse import EventSourceResponse

app = FastAPI()

app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)


@app.get("/agents.json")
def get_copilot_description():
"""Agent descriptor for the OpenBB Workspace."""
return JSONResponse(
content={
"vanilla_agent_custom_features": {
"name": "Vanilla Agent Custom Features",
"description": "A simple agent that reports its feature status.",
"image": (
"https://github.com/OpenBB-finance/copilot-for-terminal-pro/"
"assets/14093308/7da2a512-93b9-478d-90bc-b8c3dd0cabcf"
),
"endpoints": {"query": "/v1/query"},
"features": {
"streaming": True,
"widget-dashboard-select": False,
"widget-dashboard-search": False,
"deep-research": {
"label": "Deep Research",
"default": False,
"description": "Allows the copilot to do deep research",
},
"web-search": {
"label": "Web Search",
"default": True,
"description": "Allows the copilot to search the web.",
},
},
}
}
)


@app.post("/v1/query")
async def query(request: QueryRequest) -> EventSourceResponse:
"""Stream a simple greeting with feature status."""

# Check workspace_options from request payload
# workspace_options is a list like ["web-search"] or ["deep-research", "web-search"]
workspace_options = getattr(request, "workspace_options", [])

# Check which features are enabled
deep_research_enabled = "deep-research" in workspace_options
web_search_enabled = "web-search" in workspace_options

# Build the feature status message
features_msg = (
f"- Deep Research: {'✅ Enabled' if deep_research_enabled else '❌ Disabled'}\n"
f"- Web Search: {'✅ Enabled' if web_search_enabled else '❌ Disabled'}"
)

openai_messages: list[ChatCompletionMessageParam] = [
ChatCompletionSystemMessageParam(
role="system",
content=(
"You are a simple greeting agent.\n"
"Greet the user and let them know their current feature settings:\n"
f"{features_msg}\n"
"Keep your response brief and friendly."
),
)
]

for message in request.messages:
if message.role == "human":
openai_messages.append(
ChatCompletionUserMessageParam(role="user", content=message.content)
)
elif message.role == "ai" and isinstance(message.content, str):
openai_messages.append(
ChatCompletionAssistantMessageParam(
role="assistant", content=message.content
)
)

async def execution_loop() -> AsyncGenerator[MessageChunkSSE, None]:
client = openai.AsyncOpenAI()
async for event in await client.chat.completions.create(
model="gpt-4o",
messages=openai_messages,
stream=True,
):
if chunk := event.choices[0].delta.content:
yield message_chunk(chunk)

return EventSourceResponse(
content=(
event.model_dump(exclude_none=True) async for event in execution_loop()
),
media_type="text/event-stream",
)
Loading