Add Websockets example (microsoft#2291)

davorrunje · web-flow · commit 6750e67c03d4 · 2024-04-08T13:03:05.000Z
* add websockets example

* polishing

* README.me renamed to README.md

* polishing
diff --git a/samples/apps/websockets/README.md b/samples/apps/websockets/README.md
@@ -0,0 +1,20 @@
+# Using websockets with FastAPI and AutoGen
+
+## Running the example
+
+1. Navigate to the directory containing the example:
+    ```
+    cd samples/apps/websockets
+    ```
+
+2. Install the necessary dependencies:
+    ```
+    ./setup.py
+    ```
+
+3. Run the application:
+    ```
+    uvicorn application:app --reload
+    ```
+
+You should now be able to access the application in your web browser at `http://localhost:8000`.
diff --git a/samples/apps/websockets/application.py b/samples/apps/websockets/application.py
@@ -0,0 +1,182 @@
+#!/usr/bin/env python
+
+import asyncio
+import logging
+import os
+from contextlib import asynccontextmanager  # noqa: E402
+from datetime import datetime
+from typing import AsyncIterator, Dict, Iterator, List
+
+import uvicorn  # noqa: E402
+from fastapi import FastAPI  # noqa: E402
+from fastapi.responses import HTMLResponse  # noqa: E402
+from websockets.sync.client import connect as ws_connect
+
+import autogen
+from autogen.io.websockets import IOWebsockets
+
+PORT = 8000
+
+# logger = getLogger(__name__)
+logger = logging.getLogger("uvicorn")
+
+
+def _get_config_list() -> List[Dict[str, str]]:
+    """Get a list of config dictionaries with API keys for OpenAI and Azure OpenAI.
+
+    Returns:
+        List[Dict[str, str]]: A list of config dictionaries with API keys.
+
+    Example:
+        >>> _get_config_list()
+        [
+            {
+                'model': 'gpt-35-turbo-16k',
+                'api_key': '0123456789abcdef0123456789abcdef',
+                'base_url': 'https://my-deployment.openai.azure.com/',
+                'api_type': 'azure',
+                'api_version': '2024-02-15-preview',
+            },
+            {
+                'model': 'gpt-4',
+                'api_key': '0123456789abcdef0123456789abcdef',
+            },
+        ]
+    """
+    # can use both OpenAI and Azure OpenAI API keys
+    config_list = [
+        {
+            "model": "gpt-35-turbo-16k",
+            "api_key": os.environ.get("AZURE_OPENAI_API_KEY"),
+            "base_url": os.environ.get("AZURE_OPENAI_BASE_URL"),
+            "api_type": "azure",
+            "api_version": os.environ.get("AZURE_OPENAI_API_VERSION"),
+        },
+        {
+            "model": "gpt-4",
+            "api_key": os.environ.get("OPENAI_API_KEY"),
+        },
+    ]
+    # filter out configs with no API key
+    config_list = [llm_config for llm_config in config_list if llm_config["api_key"] is not None]
+
+    if not config_list:
+        raise ValueError(
+            "No API keys found. Please set either AZURE_OPENAI_API_KEY or OPENAI_API_KEY environment variable."
+        )
+
+    return config_list
+
+
+def on_connect(iostream: IOWebsockets) -> None:
+    logger.info(f"on_connect(): Connected to client using IOWebsockets {iostream}")
+
+    logger.info("on_connect(): Receiving message from client.")
+
+    # get the initial message from the client
+    initial_msg = iostream.input()
+
+    # instantiate an agent named "chatbot"
+    agent = autogen.ConversableAgent(
+        name="chatbot",
+        system_message="Complete a task given to you and reply TERMINATE when the task is done. If asked about the weather, use tool weather_forecast(city) to get the weather forecast for a city.",
+        llm_config={
+            "config_list": _get_config_list(),
+            "stream": True,
+        },
+    )
+
+    # create a UserProxyAgent instance named "user_proxy"
+    user_proxy = autogen.UserProxyAgent(
+        name="user_proxy",
+        system_message="A proxy for the user.",
+        is_termination_msg=lambda x: x.get("content", "") and x.get("content", "").rstrip().endswith("TERMINATE"),
+        human_input_mode="NEVER",
+        max_consecutive_auto_reply=10,
+        code_execution_config=False,
+    )
+
+    # register the weather_forecast function
+    def weather_forecast(city: str) -> str:
+        return f"The weather forecast for {city} at {datetime.now()} is sunny."
+
+    autogen.register_function(
+        weather_forecast, caller=agent, executor=user_proxy, description="Weather forecast for a city"
+    )
+
+    # instantiate a chat
+    logger.info(
+        f"on_connect(): Initiating chat with the agent ({agent.name}) and the user proxy ({user_proxy.name}) using the message '{initial_msg}'",
+    )
+    user_proxy.initiate_chat(  # noqa: F704
+        agent,
+        message=initial_msg,
+    )
+
+    logger.info("on_connect(): Finished the task successfully.")
+
+
+html = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>Autogen websocket test</title>
+    </head>
+    <body>
+        <h1>WebSocket Chat</h1>
+        <form action="" onsubmit="sendMessage(event)">
+            <input type="text" id="messageText" autocomplete="off" value="Write a poem about the current wearther in Paris or London, you choose."/>
+            <button>Send</button>
+        </form>
+        <ul id='messages'>
+        </ul>
+        <script>
+            var ws = new WebSocket("ws://localhost:8080/ws");
+            ws.onmessage = function(event) {
+                var messages = document.getElementById('messages')
+                var message = document.createElement('li')
+                var content = document.createTextNode(event.data)
+                message.appendChild(content)
+                messages.appendChild(message)
+            };
+            function sendMessage(event) {
+                var input = document.getElementById("messageText")
+                ws.send(input.value)
+                input.value = ''
+                event.preventDefault()
+            }
+        </script>
+    </body>
+</html>
+"""
+
+
+@asynccontextmanager
+async def run_websocket_server(app: FastAPI) -> AsyncIterator[None]:
+    with IOWebsockets.run_server_in_thread(on_connect=on_connect, port=8080) as uri:
+        logger.info(f"Websocket server started at {uri}.")
+
+        yield
+
+
+app = FastAPI(lifespan=run_websocket_server)
+
+
+@app.get("/")
+async def get() -> HTMLResponse:
+    return HTMLResponse(html)
+
+
+async def start_uvicorn() -> None:
+    config = uvicorn.Config(app)
+    server = uvicorn.Server(config)
+    try:
+        await server.serve()  # noqa: F704
+    except KeyboardInterrupt:
+        logger.info("Shutting down server")
+
+
+if __name__ == "__main__":
+    # set the log level to INFO
+    logger.setLevel("INFO")
+    asyncio.run(start_uvicorn())
diff --git a/samples/apps/websockets/setup.py b/samples/apps/websockets/setup.py
@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+
+# Equivalent to running the basj script below, but with an additional check if the files was moved:
+# cd ../../..
+# pip install -e .[websockets] fastapi uvicorn
+
+import subprocess
+from pathlib import Path
+
+repo_root = Path(__file__).parents[3]
+if not (repo_root / "setup.py").exists():
+    raise RuntimeError("This script has been moved, please run it from its original location.")
+
+print("Installing the package in editable mode, with the websockets extra, and fastapi and uvicorn...", flush=True)
+subprocess.run(["pip", "install", "-e", ".[websockets]", "fastapi", "uvicorn"], cwd=repo_root, check=True)
