pop out greptimeai keyword args

GreptimeTeam · Dec 26, 2023 · 1026d67 · 1026d67
1 parent ae8981a
commit 1026d67
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git 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
@@ -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)