From b2f2fb97b14c923b8fcde23b01c10a9a877aac83 Mon Sep 17 00:00:00 2001
From: Wenzhe Xue <wenzhe@log10.io>
Date: Wed, 21 Feb 2024 20:00:17 -0800
Subject: [PATCH] wip

---
 log10/_httpx_utils.py | 21 ++++++++++-----------
 log10/load.py         | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/log10/_httpx_utils.py b/log10/_httpx_utils.py
index efeec436..74a15333 100644
--- a/log10/_httpx_utils.py
+++ b/log10/_httpx_utils.py
@@ -6,7 +6,8 @@
 import httpx
 from dotenv import load_dotenv
 from httpx import Request, Response
-from load import global_tags, sessionID
+
+from log10.load import global_tags, sessionID
 
 
 load_dotenv()
@@ -17,8 +18,9 @@
 token = os.environ.get("LOG10_TOKEN")
 org_id = os.environ.get("LOG10_ORG_ID")
 
-transport = httpx.Transport(retries=5)
-httpx_client = httpx.Client(transport=transport)
+httpx_client = httpx.Client()
+httpx_logger = logging.getLogger("httpx")
+httpx_logger.setLevel(logging.DEBUG)
 
 
 def _try_post_request(url: str, payload: dict = {}) -> httpx.Response:
@@ -30,6 +32,7 @@ def _try_post_request(url: str, payload: dict = {}) -> httpx.Response:
     }
     res = None
     try:
+        logger.debug(f"POST {url} with headers={headers} payload: {payload}")
         res = httpx_client.post(url, headers=headers, json=payload)
         res.raise_for_status()
     except httpx.HTTPError as http_err:
@@ -39,18 +42,19 @@ def _try_post_request(url: str, payload: dict = {}) -> httpx.Response:
                 + "\nSee https://github.com/log10-io/log10#%EF%B8%8F-setup for details"
             )
         else:
-            logger.error(f"Log10: failed with error: {http_err}")
+            logger.error(f"Failed with error: {http_err}")
     except Exception as err:
-        logger.error(f"Log10: failed to insert in log10: {payload} with error {err}")
+        logger.error(f"Failed to insert in log10: {payload} with error {err}")
 
 
 async def get_completion_id(request: Request):
     completion_url = "/api/completions"
     res = _try_post_request(url=f"{base_url}{completion_url}")
+    import ipdb; ipdb.set_trace()
     try:
         completion_id = res.json().get("completionID")
     except Exception as e:
-        logger.error(f"LOG10: failed to get completion ID. Error: {e}. Skipping completion recording.")
+        logger.error(f"Failed to get completion ID. Error: {e}. Skipping completion recording.")
     else:
         request.headers["x-log10-completion-id"] = completion_id
 
@@ -179,8 +183,3 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
                 extensions=response.extensions,
                 request=request,
             )
-
-async def 
-
-
-
diff --git a/log10/load.py b/log10/load.py
index 018d79e9..a00fbbc0 100644
--- a/log10/load.py
+++ b/log10/load.py
@@ -642,6 +642,24 @@ def log10(module, DEBUG_=False, USE_ASYNC_=True):
             attr = module.resources.chat.completions.Completions
             method = getattr(attr, "create")
             setattr(attr, "create", intercepting_decorator(method))
+
+            # support for async completions
+            # patch module.AsyncOpenAI.__init__ to new_init
+            origin_init = module.AsyncOpenAI.__init__
+            def new_init(self, *args, **kwargs):
+                logger.debug("LOG10: patching AsyncOpenAI.__init__")
+                import httpx
+
+                from log10._httpx_utils import get_completion_id
+                event_hooks = {
+                    "request": [get_completion_id],
+                }
+                async_httpx_client = httpx.AsyncClient(
+                    event_hooks=event_hooks)
+                kwargs["http_client"] = async_httpx_client
+                origin_init(self, *args, **kwargs)
+            module.AsyncOpenAI.__init__ = new_init
+
         else:
             attr = module.api_resources.completion.Completion
             method = getattr(attr, "create")