From 1026d679ae0fa99c35fb1e4ed0ff8173da7d2314 Mon Sep 17 00:00:00 2001 From: yuanbohan Date: Tue, 26 Dec 2023 11:22:09 +0800 Subject: [PATCH] pop out greptimeai keyword args --- src/greptimeai/extractor/openai_extractor/__init__.py | 9 +++++++++ src/greptimeai/patcher/openai_patcher/base.py | 2 ++ 2 files changed, 11 insertions(+) diff --git a/src/greptimeai/extractor/openai_extractor/__init__.py b/src/greptimeai/extractor/openai_extractor/__init__.py index b77847b..c31f8a6 100644 --- a/src/greptimeai/extractor/openai_extractor/__init__.py +++ b/src/greptimeai/extractor/openai_extractor/__init__.py @@ -66,6 +66,15 @@ def extract_req_tokens(**kwargs) -> Optional[Union[str, list]]: logger.warning(f"Failed to extract req tokens from {kwargs=}") return None + @staticmethod + def pop_out_keyword_args(kwargs: Dict[str, Any]): + """ + pop out the keyword args which are not supported by openai + """ + kwargs.pop(_GREPTIMEAI_USER_KEY, None) + kwargs.pop(_GREPTIMEAI_TRACE_KEY, None) + kwargs.pop(_GREPTIMEAI_SPAN_KEY, None) + @staticmethod def update_trace_info(kwargs: Dict[str, Any], trace_id: str, span_id: str): attrs = { diff --git a/src/greptimeai/patcher/openai_patcher/base.py b/src/greptimeai/patcher/openai_patcher/base.py index feb3c05..880ca56 100644 --- a/src/greptimeai/patcher/openai_patcher/base.py +++ b/src/greptimeai/patcher/openai_patcher/base.py @@ -96,7 +96,9 @@ def _pre_patch( span_attrs=extraction.span_attributes, event_attrs=extraction.event_attributes, ) + OpenaiExtractor.update_trace_info(kwargs, trace_id, span_id) + OpenaiExtractor.pop_out_keyword_args(kwargs) start = time.time() return (extraction, span_id, start, kwargs)