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

Stoli/feat/connection handling #568

Merged

Conversation

apostolos-geyer
Copy link
Contributor

Updates regarding my comments in #567.

The changes are as follows:

Addition of a LlamaParse.aclient property.

The aclient property handles setting of the base_url, timeout, as well as the Authorization header on the client, such that the logic does not have to be repeated in each method that uses the client.

It will use custom_client in cases where one is provided, so a user could still provide a subclass of httpx.AsyncClient or one pre-initialized with event hooks.

Methods no longer use the client_context async context manager.

Using httpx.Client / httpx.AsyncClient as a context manager results in the cleanup of the client and all connection persistence once the context manager block is exited. (See: https://github.com/encode/httpx/blob/c7c13f18a5af4c64c649881b2fe8dbd72a519c32/httpx/_client.py#L2014).

It is better to use it as a context manager for one-offs, wherein you want to send a few requests and keep the session alive for those few requests, before cleaning it up.

For LlamaParse's use case, it makes more sense to keep a single client instance so that connections are kept alive. For this reason, I've updated the methods that make requests to use the client from the aclient property.

I have preserved the client_context method for interface compatibility purposes in case any users or other libraries have code that depends on its presence.

You can see the difference in behaviour by running this simple example in a notebook:

import logging

import httpx
from llama_parse import LlamaParse

logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(message)s',  # Include timestamp, log level, and message
    force=True
)
logging.getLogger('httpx').setLevel(logging.DEBUG)
logging.getLogger('httpcore').setLevel(logging.DEBUG)

test = "tests/test_files/attention_is_all_you_need.pdf"

parse = LlamaParse()

resps: list[httpx.Response] = []
async def save_resps(resp: httpx.Response) -> httpx.Response:
    resps.append(resp)
    return resp

parse.aclient.event_hooks['response'] = [save_resps]

for i,t in enumerate([test] * 5):
    # send the request 5 times.. i know it accepts a list, but it seems the underlying effect is the same and this is a bit easier to differentiate

    print(f'{"-"*50}', flush=True)
    print(f'request {i=}', flush=True)
    _ = await parse.aget_json(t)

ssl_objs = set([r.extensions["network_stream"].get_extra_info("ssl_object") for r in resps])

print(f'{len(ssl_objs)=}')

You'll notice that the TCP and TLS Handshakes occur only once on the first request, and the connection is reused.

This is further illustrated by len(ssl_objs)=1. What this means is that all requests are using the same underlying stream.

Everything else:

The below are not relevant to the primary changes made, but also included in the PR. I can remove them if necessary, as I completely understand a preference for keeping change-sets minimal and focused in the interest of keeping branch history clean.

  • _create_job preserves exception context if the request is not successful:
    Previously, _create_job would raise an Exception if response.is_success was false (i.e. if the status code is not 2XX). I changed this to use the raise_for_status method of httpx.Response, and then raise Exception(...) from ( the exception raised by raise_for_status )... This is entirely unimportant to 99% of developers but does assist in debugging as it means one can access the raw Request and Response objects when an exception is caught, as raise ... from ... will store the RHS of from on the __cause__ attribute of an Exception.

  • Imports are now sorted: Honestly I just forgot to turn off ruff and my global configuration sorts imports on save. Hope this isn't too much of a bother.

P.S: Wouldn't be a bad idea to have a develop branch.... Feels kinda unnerving making a PR right to master.

Thanks for the opportunity to contribute, you guys are doing great work.

@logan-markewich logan-markewich self-assigned this Dec 27, 2024
@logan-markewich
Copy link
Contributor

LGTM, tests pass. Thanks!

@logan-markewich logan-markewich merged commit 530241d into run-llama:main Dec 27, 2024
9 checks passed
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