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)