From 0ddd72081efc6bf7afb8e6408414f375481c5097 Mon Sep 17 00:00:00 2001 From: Rohit Kadhe Date: Wed, 11 Sep 2024 12:14:38 -0600 Subject: [PATCH 1/3] support genai and also add token reporting --- src/examples/langchain_example/__init__.py | 2 + .../langchain_google_genai.py | 21 +++++++ .../langchain/instrumentation.py | 9 +-- .../langchain_core/instrumentation.py | 14 +++++ .../instrumentation/langchain_core/patch.py | 55 +++++++++++++------ src/langtrace_python_sdk/langtrace.py | 6 +- src/run_example.py | 2 +- 7 files changed, 81 insertions(+), 28 deletions(-) create mode 100644 src/examples/langchain_example/langchain_google_genai.py diff --git a/src/examples/langchain_example/__init__.py b/src/examples/langchain_example/__init__.py index f421b6b1..ceaee24c 100644 --- a/src/examples/langchain_example/__init__.py +++ b/src/examples/langchain_example/__init__.py @@ -1,3 +1,4 @@ +from examples.langchain_example.langchain_google_genai import basic_google_genai from .basic import basic_app, rag, load_and_split from langtrace_python_sdk import with_langtrace_root_span @@ -12,6 +13,7 @@ def run(self): rag() load_and_split() basic_graph_tools() + basic_google_genai() class GroqRunner: diff --git a/src/examples/langchain_example/langchain_google_genai.py b/src/examples/langchain_example/langchain_google_genai.py new file mode 100644 index 00000000..78e9fa52 --- /dev/null +++ b/src/examples/langchain_example/langchain_google_genai.py @@ -0,0 +1,21 @@ +from langchain_core.messages import HumanMessage +from langchain_google_genai import ChatGoogleGenerativeAI +from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span + + +@with_langtrace_root_span("basic_google_genai") +def basic_google_genai(): + llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash") + # example + message = HumanMessage( + content=[ + { + "type": "text", + "text": "What's in this image?", + }, + ] + ) + message_image = HumanMessage(content="https://picsum.photos/seed/picsum/200/300") + + res = llm.invoke([message, message_image]) + # print(res) diff --git a/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py b/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py index 4cc4fe8f..7d01906e 100644 --- a/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py +++ b/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py @@ -81,21 +81,16 @@ def _instrument(self, **kwargs): tracer_provider = kwargs.get("tracer_provider") tracer = get_tracer(__name__, "", tracer_provider) version = importlib.metadata.version("langchain") - wrap_function_wrapper( "langchain.agents.agent", "RunnableAgent.plan", - generic_patch( - "RunnableAgent.plan", "plan", tracer, version, True, True - ), + generic_patch("RunnableAgent.plan", "plan", tracer, version, True, True), ) wrap_function_wrapper( "langchain.agents.agent", "RunnableAgent.aplan", - generic_patch( - "RunnableAgent.aplan", "plan", tracer, version, True, True - ), + generic_patch("RunnableAgent.aplan", "plan", tracer, version, True, True), ) # modules_to_patch = [] diff --git a/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py b/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py index e10a6510..a7567ff0 100644 --- a/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py +++ b/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py @@ -134,6 +134,20 @@ def _instrument(self, **kwargs): ] modules_to_patch = [ + ( + "langchain_core.language_models.chat_models", + "chatmodel", + generic_patch, + True, + True, + ), + ( + "langchain_core.language_models.base", + "language_model", + generic_patch, + True, + True, + ), ("langchain_core.retrievers", "retriever", generic_patch, True, True), ("langchain_core.prompts.chat", "prompt", generic_patch, True, True), ( diff --git a/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py b/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py index 765c738e..ce93c4c5 100644 --- a/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py +++ b/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py @@ -57,7 +57,24 @@ def traced_method(wrapped, instance, args, kwargs): "langtrace.service.version": version, "langtrace.version": v(LANGTRACE_SDK_NAME), "langchain.task.name": task, - **(extra_attributes if extra_attributes is not None else {}), + "gen_ai.request.model": ( + instance.model if hasattr(instance, "model") else None + ), + SpanAttributes.LLM_REQUEST_MAX_TOKENS: ( + instance.max_output_tokens + if hasattr(instance, "max_output_tokens") + else None + ), + SpanAttributes.LLM_TOP_K: ( + instance.top_k if hasattr(instance, "top_k") else None + ), + SpanAttributes.LLM_REQUEST_TOP_P: ( + instance.top_p if hasattr(instance, "top_p") else None + ), + SpanAttributes.LLM_REQUEST_TEMPERATURE: ( + instance.temperature if hasattr(instance, "temperature") else None + ), + **(extra_attributes if extra_attributes is not None else {}), # type: ignore } if trace_input and len(args) > 0: @@ -79,21 +96,17 @@ def traced_method(wrapped, instance, args, kwargs): try: # Attempt to call the original method result = wrapped(*args, **kwargs) - if trace_output: span.set_attribute("langchain.outputs", to_json_string(result)) - if hasattr(result, 'usage'): - prompt_tokens = result.usage.prompt_tokens - completion_tokens = result.usage.completion_tokens - span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens) - span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens) - elif hasattr(result, 'generations') and len(result.generations) > 0 and len(result.generations[0]) > 0 and hasattr(result.generations[0][0], 'text') and isinstance(result.generations[0][0].text, str): - span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, instance.get_num_tokens(result.generations[0][0].text)) - elif len(args) > 0 and len(args[0]) > 0 and not hasattr(args[0][0], 'text') and hasattr(instance, 'get_num_tokens'): - span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0])) - elif len(args) > 0 and len(args[0]) > 0 and hasattr(args[0][0], 'text') and isinstance(args[0][0].text, str) and hasattr(instance, 'get_num_tokens'): - span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0].text)) - + if hasattr(result, "usage_metadata"): + span.set_attribute( + SpanAttributes.LLM_USAGE_PROMPT_TOKENS, + result.usage_metadata["input_tokens"], + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, + result.usage_metadata["output_tokens"], + ) span.set_status(StatusCode.OK) return result except Exception as err: @@ -208,9 +221,17 @@ def clean_empty(d): if not isinstance(d, (dict, list, tuple)): return d if isinstance(d, tuple): - return tuple(val for val in (clean_empty(val) for val in d) if val != () and val is not None) + return tuple( + val + for val in (clean_empty(val) for val in d) + if val != () and val is not None + ) if isinstance(d, list): - return [val for val in (clean_empty(val) for val in d) if val != [] and val is not None] + return [ + val + for val in (clean_empty(val) for val in d) + if val != [] and val is not None + ] result = {} for k, val in d.items(): if isinstance(val, dict): @@ -226,7 +247,7 @@ def clean_empty(d): result[k] = val.strip() elif isinstance(val, object): # some langchain objects have a text attribute - val = getattr(val, 'text', None) + val = getattr(val, "text", None) if val is not None and val.strip() != "": result[k] = val.strip() return result diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index 01c6417d..f5682ca1 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/src/langtrace_python_sdk/langtrace.py @@ -128,8 +128,8 @@ def init( "embedchain": EmbedchainInstrumentation(), "qdrant-client": QdrantInstrumentation(), "langchain": LangchainInstrumentation(), - "langchain-core": LangchainCoreInstrumentation(), - "langchain-community": LangchainCommunityInstrumentation(), + "langchain_core": LangchainCoreInstrumentation(), + "langchain_community": LangchainCommunityInstrumentation(), "langgraph": LanggraphInstrumentation(), "anthropic": AnthropicInstrumentation(), "cohere": CohereInstrumentation(), @@ -190,7 +190,7 @@ def init( def init_instrumentations( - disable_instrumentations: DisableInstrumentations, all_instrumentations: dict + disable_instrumentations: DisableInstrumentations | None, all_instrumentations: dict ): if disable_instrumentations is None: for idx, (name, v) in enumerate(all_instrumentations.items()): diff --git a/src/run_example.py b/src/run_example.py index cddd1e42..2393c385 100644 --- a/src/run_example.py +++ b/src/run_example.py @@ -2,7 +2,7 @@ ENABLED_EXAMPLES = { "anthropic": False, - "azureopenai": True, + "azureopenai": False, "chroma": False, "cohere": False, "fastapi": False, From fcaf4b0c2075ad4fd8b58eda243dfff444e91bb8 Mon Sep 17 00:00:00 2001 From: Karthik Kalyanaraman Date: Wed, 11 Sep 2024 15:03:54 -0500 Subject: [PATCH 2/3] rename example file --- ...gchain_google_genai.py => langchain_google_example.py} | 8 ++++++++ 1 file changed, 8 insertions(+) rename src/examples/langchain_example/{langchain_google_genai.py => langchain_google_example.py} (80%) diff --git a/src/examples/langchain_example/langchain_google_genai.py b/src/examples/langchain_example/langchain_google_example.py similarity index 80% rename from src/examples/langchain_example/langchain_google_genai.py rename to src/examples/langchain_example/langchain_google_example.py index 78e9fa52..fc92c8b9 100644 --- a/src/examples/langchain_example/langchain_google_genai.py +++ b/src/examples/langchain_example/langchain_google_example.py @@ -1,7 +1,12 @@ from langchain_core.messages import HumanMessage from langchain_google_genai import ChatGoogleGenerativeAI from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span +from dotenv import find_dotenv, load_dotenv +from langtrace_python_sdk import langtrace +_ = load_dotenv(find_dotenv()) + +langtrace.init() @with_langtrace_root_span("basic_google_genai") def basic_google_genai(): @@ -19,3 +24,6 @@ def basic_google_genai(): res = llm.invoke([message, message_image]) # print(res) + + +basic_google_genai() From 94bead33c8d29746996e6b49ad4b1b922a05856c Mon Sep 17 00:00:00 2001 From: Karthik Kalyanaraman Date: Wed, 11 Sep 2024 15:09:54 -0500 Subject: [PATCH 3/3] fix --- src/langtrace_python_sdk/langtrace.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index f5682ca1..d0aacddd 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/src/langtrace_python_sdk/langtrace.py @@ -190,7 +190,8 @@ def init( def init_instrumentations( - disable_instrumentations: DisableInstrumentations | None, all_instrumentations: dict + disable_instrumentations: Optional[DisableInstrumentations], + all_instrumentations: dict ): if disable_instrumentations is None: for idx, (name, v) in enumerate(all_instrumentations.items()):