60 lines
1.4 KiB
Python
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)
|