Updated the wrapper method format

monocle2ai · Jul 26, 2024 · 19b8559 · 19b8559
1 parent 712eb99
commit 19b8559
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 55 deletions.
diff --git a/src/monocle_apptrace/utils.py b/src/monocle_apptrace/utils.py
@@ -60,16 +60,13 @@ def load_wrapper_from_config(config_file_path:str, module_name:str=None):
         wrapper_methods = json_data["wrapper_methods"]
         for wrapper_method in wrapper_methods:
             wrapper_method["wrapper"] = get_wrapper_method(
-                                wrapper_method["wrapper"])
-            if "span_name_getter" in wrapper_method :
+                wrapper_method["wrapper_package"], wrapper_method["wrapper_method"])
+            if "span_name_getter_method" in wrapper_method :
                 wrapper_method["span_name_getter"] = get_wrapper_method(
-                                        wrapper_method["span_name_getter"])
+                    wrapper_method["span_name_getter_package"],
+                    wrapper_method["span_name_getter_method"])
         return wrapper_methods
 
-def get_wrapper_method(method_signature: str):
-    method_tokens = method_signature.split(".")
-    module_name = ""
-    for i in range(len(method_tokens) -1):
-        module_name += "." + method_tokens[i]
-    wrapper_module = import_module("monocle_apptrace" + module_name)
-    return getattr(wrapper_module, method_tokens[len(method_tokens)-1])
+def get_wrapper_method(package_name: str, method_name: str):
+    wrapper_module = import_module("monocle_apptrace." + package_name)
+    return getattr(wrapper_module, method_name)
diff --git a/src/monocle_apptrace/wrapper_config/haystack_methods.json b/src/monocle_apptrace/wrapper_config/haystack_methods.json
@@ -1,22 +1,25 @@
 {
-   "wrapper_methods" :  [
-        {
-            "package": "haystack.components.generators.openai",
-            "object": "OpenAIGenerator",
-            "method": "run",
-            "wrapper": "haystack.wrap_openai.wrap_openai"
-        },
-        {
-            "package": "haystack.components.generators.chat.openai",
-            "object": "OpenAIChatGenerator",
-            "method": "run",
-            "wrapper": "haystack.wrap_openai.wrap_openai"
-        },
-        {
-            "package": "haystack.core.pipeline.pipeline",
-            "object": "Pipeline",
-            "method": "run",
-            "wrapper": "haystack.wrap_pipeline.wrap"
-        }
-    ] 
+"wrapper_methods" :  [
+    {
+        "package": "haystack.components.generators.openai",
+        "object": "OpenAIGenerator",
+        "method": "run",
+        "wrapper_package": "haystack.wrap_openai",
+        "wrapper_method": "wrap_openai"
+    },
+    {
+        "package": "haystack.components.generators.chat.openai",
+        "object": "OpenAIChatGenerator",
+        "method": "run",
+        "wrapper_package": "haystack.wrap_openai",
+        "wrapper_method": "wrap_openai"
+    },
+    {
+        "package": "haystack.core.pipeline.pipeline",
+        "object": "Pipeline",
+        "method": "run",
+        "wrapper_package": "haystack.wrap_pipeline",
+        "wrapper_method": "wrap"
+    }
+] 
 }
diff --git a/src/monocle_apptrace/wrapper_config/lang_chain_methods.json b/src/monocle_apptrace/wrapper_config/lang_chain_methods.json
@@ -4,89 +4,103 @@
         "package": "langchain.prompts.base",
         "object": "BasePromptTemplate",
         "method": "invoke",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "langchain.prompts.base",
         "object": "BasePromptTemplate",
         "method": "ainvoke",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "langchain.chat_models.base",
         "object": "BaseChatModel",
         "method": "invoke",
-        "wrapper": "wrap_common.llm_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper"
     },
     {
         "package": "langchain.chat_models.base",
         "object": "BaseChatModel",
         "method": "ainvoke",
-        "wrapper": "wrap_common.allm_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "allm_wrapper"
     },
     {
         "package": "langchain_core.language_models.llms",
         "object": "LLM",
         "method": "_generate",
-        "wrapper": "wrap_common.llm_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper"
     },
     {
         "package": "langchain_core.language_models.llms",
         "object": "LLM",
         "method": "_agenerate",
-        "wrapper": "wrap_common.llm_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper"
     },
     {
         "package": "langchain_core.retrievers",
         "object": "BaseRetriever",
         "method": "invoke",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "langchain_core.retrievers",
         "object": "BaseRetriever",
         "method": "ainvoke",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "langchain.schema",
         "object": "BaseOutputParser",
         "method": "invoke",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "langchain.schema",
         "object": "BaseOutputParser",
         "method": "ainvoke",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "langchain.schema.runnable",
         "object": "RunnableSequence",
         "method": "invoke",
         "span_name": "langchain.workflow",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "langchain.schema.runnable",
         "object": "RunnableSequence",
         "method": "ainvoke",
         "span_name": "langchain.workflow",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "langchain.schema.runnable",
         "object": "RunnableParallel",
         "method": "invoke",
         "span_name": "langchain.workflow",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "langchain.schema.runnable",
         "object": "RunnableParallel",
         "method": "ainvoke",
         "span_name": "langchain.workflow",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     }
 ]
-}
+}
diff --git a/src/monocle_apptrace/wrapper_config/llama_index_methods.json b/src/monocle_apptrace/wrapper_config/llama_index_methods.json
@@ -5,57 +5,66 @@
         "object": "BaseRetriever",
         "method": "retrieve",
         "span_name": "llamaindex.retrieve",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "llama_index.core.indices.base_retriever",
         "object": "BaseRetriever",
         "method": "aretrieve",
         "span_name": "llamaindex.retrieve",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "llama_index.core.base.base_query_engine",
         "object": "BaseQueryEngine",
         "method": "query",
         "span_name": "llamaindex.query",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "llama_index.core.base.base_query_engine",
         "object": "BaseQueryEngine",
         "method": "aquery",
         "span_name": "llamaindex.query",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "llama_index.core.llms.custom",
         "object": "CustomLLM",
         "method": "chat",
         "span_name": "llamaindex.llmchat",
-        "wrapper": "wrap_common.task_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
     },
     {
         "package": "llama_index.core.llms.custom",
         "object": "CustomLLM",
         "method": "achat",
         "span_name": "llamaindex.llmchat",
-        "wrapper": "wrap_common.atask_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
     },
     {
         "package": "llama_index.llms.openai.base",
         "object": "OpenAI",
         "method": "chat",
         "span_name": "llamaindex.openai",
-        "span_name_getter" : "llamaindex.get_llm_span_name_for_openai",
-        "wrapper": "wrap_common.llm_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper",
+        "span_name_getter_package" : "llamaindex",
+        "span_name_getter_mothod" : "get_llm_span_name_for_openai"
     },
     {
         "package": "llama_index.llms.openai.base",
         "object": "OpenAI",
         "method": "achat",
         "span_name": "llamaindex.openai",
-        "wrapper": "wrap_common.allm_wrapper"
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "allm_wrapper"
     }
 ]
-}
+}