60 lines
1.4 KiB
Python

from __future__ import annotations
from typing import Final
import socket
SERVER_HOST: Final = "127.0.0.1"
SERVER_PORT: Final = 9000
SHARED_SECRET: Final = "change-me"
from construct import ConstructError, Int32ul, PascalString, PrefixedArray, Struct as CStruct
MessageInChat = CStruct(
"role" / PascalString(Int32ul, "utf8"),
"content" / PascalString(Int32ul, "utf8"),
)
IncomingChatBatch = CStruct(
"messages" / PrefixedArray(Int32ul, MessageInChat),
)
from secret_stream_socket import SecretStreamSocket, wrap_connection_socket
def connect_to_server(
host: str = SERVER_HOST,
port: int = SERVER_PORT,
shared_secret: str = SHARED_SECRET,
) -> SecretStreamSocket:
sock = socket.create_connection((host, port))
return wrap_connection_socket(sock, shared_secret)
def send_message(connection: SecretStreamSocket, role: str, content: str) -> None:
payload = IncomingChatBatch.build(
{
"messages": [
{
"role": role,
"content": content,
}
]
}
)
connection.send_frame(payload)
def close_connection(connection: SecretStreamSocket) -> None:
connection.close()
if __name__ == "__main__":
connection = connect_to_server()
try:
msg = "\n".join([f"hello {i} from F.py" for i in range(1000)])
send_message(connection, "user", msg)
finally:
close_connection(connection)