Add a way to stream output to a UI via sockets

autogen/oai/client.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 import os
 import sys
 from typing import Any, List, Optional, Dict, Callable, Tuple, Union
@@ -16,6 +15,8 @@
 from autogen.token_count_utils import count_token
 from autogen._pydantic import model_dump
 
+import socket_connection_singleton
+
 TOOL_ENABLED = False
 try:
     import openai
@@ -488,6 +489,8 @@ def _completions_create(self, client: OpenAI, params: Dict[str, Any]) -> ChatCom
 
                         # If content is present, print it to the terminal and update response variables
                         if content is not None:
+                            if socket_connection_singleton.clientsocket is not None:
+                                socket_connection_singleton.clientsocket.send(bytes(str(content), "utf-8"))
                             print(content, end="", flush=True)
                             response_contents[choice.index] += content
                             completion_tokens += 1

autogen/oai/socket_connection_singleton.py

@@ -0,0 +1,30 @@
+#This is a singleton.  You use it to pass variables around from script to script, on any level.
+#Import this script into your top script and define your socket connection, then import this script's variables 
+#into the file where you want to use them.
+
+#For example:
+
+#main.py:
+#import socket
+#import socket_connection_singleton
+
+#s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+#s.bind((socket.gethostname(), 1235))
+#s.listen(5)
+#logging.info("Waiting for connection...")
+#clientsocket, address = s.accept()
+#socket_connection_singleton.s = s
+#socket_connection_singleton.clientsocket = clientsocket
+
+
+#client.py:
+
+#import socket
+#import socket_connection_singleton
+
+#clientsocket = socket_connection_singleton.clientsocket
+#clientsocket.send(bytes("Hello, World!", "utf-8"))
+
+s = None
+clientsocket = None
+address = None