From 19b855985b0d98e042db98d2a63bb73cdc175582 Mon Sep 17 00:00:00 2001 From: Prasad Mujumdar Date: Thu, 25 Jul 2024 19:21:26 -0700 Subject: [PATCH] Updated the wrapper method format --- src/monocle_apptrace/utils.py | 17 +++---- .../wrapper_config/haystack_methods.json | 43 +++++++++--------- .../wrapper_config/lang_chain_methods.json | 44 ++++++++++++------- .../wrapper_config/llama_index_methods.json | 29 +++++++----- 4 files changed, 78 insertions(+), 55 deletions(-) diff --git a/src/monocle_apptrace/utils.py b/src/monocle_apptrace/utils.py index 9cd1059..bc8f0af 100644 --- 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 index f77d83f..e02b21c 100644 --- 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 index d1b45c8..8de8764 100644 --- 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" } ] -} \ No newline at end of file +} diff --git a/src/monocle_apptrace/wrapper_config/llama_index_methods.json b/src/monocle_apptrace/wrapper_config/llama_index_methods.json index 58ead74..000b540 100644 --- 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" } ] -} \ No newline at end of file +}