Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add streaming support (both steps and tokens) to the Python REST client + pytest #1701

Merged
merged 4 commits into from
Sep 3, 2024

Conversation

cpacker
Copy link
Collaborator

@cpacker cpacker commented Aug 30, 2024

This PR adds support for streaming to the Python REST client. See the new test code for an example:

def test_streaming_send_message(client: Union[LocalClient, RESTClient], agent: AgentState):
    if isinstance(client, LocalClient):
        pytest.skip("Skipping test_streaming_send_message because LocalClient does not support streaming")
    assert isinstance(client, RESTClient), client

    # First, try streaming just steps
    # TODO

    # Next, try streaming both steps and tokens
    response = client.send_message(
        agent_id=agent.id,
        message="This is a test. Repeat after me: 'banana'",
        role="user",
        stream_steps=True,
        stream_tokens=True,
    )

    # Some manual checks to run
    # 1. Check that there were inner thoughts
    inner_thoughts_exist = False
    # 2. Check that the agent runs `send_message`
    send_message_ran = False
    # 3. Check that we get all the start/stop/end tokens we want
    #    This includes all of the MessageStreamStatus enums
    done_gen = False
    done_step = False
    done = False

    # print(response)
    assert response, "Sending message failed"
    for chunk in response:
        assert isinstance(chunk, MemGPTStreamingResponse)
        if isinstance(chunk, InternalMonologue) and chunk.internal_monologue and chunk.internal_monologue != "":
            inner_thoughts_exist = True
        if isinstance(chunk, FunctionCallMessage) and chunk.function_call and chunk.function_call.name == "send_message":
            send_message_ran = True
        if isinstance(chunk, MessageStreamStatus):
            if chunk == MessageStreamStatus.done:
                assert not done, "Message stream already done"
                done = True
            elif chunk == MessageStreamStatus.done_step:
                assert not done_step, "Message stream already done step"
                done_step = True
            elif chunk == MessageStreamStatus.done_generation:
                assert not done_gen, "Message stream already done generation"
                done_gen = True
  • Add support for streaming steps (stream_steps) to Python REST client
  • Add support for streaming tokens (stream_tokens) to Python REST client

Note: punted to future PR

  • Add support for streaming steps (stream_steps) to Python local client
  • Add support for streaming tokens (stream_tokens) to Python local client

Add streaming tests to the client:

pytest tests/test_client.py::test_streaming_send_message

@cpacker cpacker changed the title feat: add streaming support (both steps and tokens) to the Python client + pytest feat: add streaming support (both steps and tokens) to the Python REST client + pytest Sep 3, 2024
Copy link
Collaborator

@sarahwooders sarahwooders left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm!

@cpacker cpacker merged commit 341d2e1 into main Sep 3, 2024
15 checks passed
@cpacker cpacker deleted the add-streaming-support-to-client branch September 3, 2024 23:57
mattzh72 pushed a commit that referenced this pull request Oct 9, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants