Merge pull request #81 from OpenBB-finance/feat/custom-feats

Author: DidierRLopes

37-vanilla-agent-custom-features/README.md

@@ -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

37-vanilla-agent-custom-features/tests/__init__.py

@@ -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

37-vanilla-agent-custom-features/tests/test_agent.py

@@ -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",
+    )