Requests made in jupyter notebooks are not beamed

[motin]

Steps to reproduce:

1. OPENAI_API_KEY="..." tokmon --no_json --beam localhost:9000 jupyter-notebook
2. Create a new notebook
3. Paste the contents of https://github.com/yagil/tokmon/blob/main/tests/python_example.py into a cell
4. Run the cell

Expected results: The request to the OpenAI API is captured and shows up in the token-beam UI

Actual results: The request completes successfully in the notebook but nothing shows up in the boken-beam UI.

Observations:

• tokmon --beam localhost:9000 python python_example.py via the cli works as expected (shows up properly in the token-beam UI), so the setup is working under ordinary circumstances.
• os.environ['HTTP_PROXY'] contains 'http://localhost:7878' when running in a notebook, so the env vars are getting set properly

[yagil]

@motin thanks for the bug report. Does it work with vanilla tokmon (without the beam UI)?

OPENAI_API_KEY="..." tokmon jupyter-notebook

[motin]

It does, if I gracefully shut down the Jupyter server (not on keyboard interrupt):

tokmon cost report:
================================================================================
Monitored invocation: jupyter-notebook notebooks/
Models: ['gpt-3.5-turbo-0301']
Total Usage: {'total_prompt_tokens': 20, 'total_completion_tokens': 13, 'total_tokens': 33}
Pricing: {'gpt-3.5-turbo-0301': {'prompt_cost': 0.002, 'completion_cost': 0.002, 'per_tokens': 1000}}
Total Cost: $0.000066
================================================================================

Writing cost summary to JSON file: /tmp/tokmon_usage_summary_1686771165.json (run with --no_json to disable this behavior)

But no luck in beam mode...