Add flexiblity to pass span through additional attributes

alizenhom · alizenhom · commit 38878b510d26 · 2024-08-05T17:00:43.000+03:00
diff --git a/src/langtrace_python_sdk/instrumentation/anthropic/patch.py b/src/langtrace_python_sdk/instrumentation/anthropic/patch.py
@@ -23,6 +23,7 @@
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
+    get_span_name,
     is_streaming,
     set_event_completion,
     set_event_completion_chunk,
@@ -64,7 +65,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         span = tracer.start_span(
-            APIS["MESSAGES_CREATE"]["METHOD"], kind=SpanKind.CLIENT
+            name=get_span_name(APIS["MESSAGES_CREATE"]["METHOD"]), kind=SpanKind.CLIENT
         )
         for field, value in attributes.model_dump(by_alias=True).items():
             set_span_attribute(span, field, value)
diff --git a/src/langtrace_python_sdk/instrumentation/chroma/patch.py b/src/langtrace_python_sdk/instrumentation/chroma/patch.py
@@ -16,6 +16,7 @@
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
 from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.llm import get_span_name
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
@@ -60,7 +61,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = DatabaseSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            api["METHOD"],
+            name=get_span_name(api["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/cohere/patch.py b/src/langtrace_python_sdk/instrumentation/cohere/patch.py
@@ -21,6 +21,7 @@
     get_llm_request_attributes,
     get_extra_attributes,
     get_llm_url,
+    get_span_name,
     set_event_completion,
     set_event_completion_chunk,
     set_usage_attributes,
@@ -57,7 +58,9 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = LLMSpanAttributes(**span_attributes)
 
-        span = tracer.start_span(APIS["RERANK"]["METHOD"], kind=SpanKind.CLIENT)
+        span = tracer.start_span(
+            name=get_span_name(APIS["RERANK"]["METHOD"]), kind=SpanKind.CLIENT
+        )
         for field, value in attributes.model_dump(by_alias=True).items():
             set_span_attribute(span, field, value)
         try:
@@ -137,7 +140,10 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = LLMSpanAttributes(**span_attributes)
 
-        span = tracer.start_span(APIS["EMBED"]["METHOD"], kind=SpanKind.CLIENT)
+        span = tracer.start_span(
+            name=get_span_name(APIS["EMBED"]["METHOD"]),
+            kind=SpanKind.CLIENT,
+        )
         for field, value in attributes.model_dump(by_alias=True).items():
             set_span_attribute(span, field, value)
         try:
@@ -225,7 +231,9 @@ def traced_method(wrapped, instance, args, kwargs):
             # stringify the list of objects
             attributes.llm_tool_results = json.dumps(kwargs.get("tool_results"))
 
-        span = tracer.start_span(APIS["CHAT_CREATE"]["METHOD"], kind=SpanKind.CLIENT)
+        span = tracer.start_span(
+            name=get_span_name(APIS["CHAT_CREATE"]["METHOD"]), kind=SpanKind.CLIENT
+        )
 
         # Set the attributes on the span
         for field, value in attributes.model_dump(by_alias=True).items():
@@ -391,7 +399,9 @@ def traced_method(wrapped, instance, args, kwargs):
             # stringify the list of objects
             attributes.llm_tool_results = json.dumps(kwargs.get("tool_results"))
 
-        span = tracer.start_span(APIS["CHAT_STREAM"]["METHOD"], kind=SpanKind.CLIENT)
+        span = tracer.start_span(
+            name=get_span_name(APIS["CHAT_STREAM"]["METHOD"]), kind=SpanKind.CLIENT
+        )
         for field, value in attributes.model_dump(by_alias=True).items():
             set_span_attribute(span, field, value)
         try:
diff --git a/src/langtrace_python_sdk/instrumentation/crewai/patch.py b/src/langtrace_python_sdk/instrumentation/crewai/patch.py
@@ -2,6 +2,7 @@
 from importlib_metadata import version as v
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.llm import get_span_name
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -143,7 +144,9 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(operation_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            get_span_name(operation_name), kind=SpanKind.CLIENT
+        ) as span:
             _set_input_attributes(span, kwargs, attributes)
 
             try:
diff --git a/src/langtrace_python_sdk/instrumentation/gemini/patch.py b/src/langtrace_python_sdk/instrumentation/gemini/patch.py
@@ -10,6 +10,7 @@
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
+    get_span_name,
     is_streaming,
     set_event_completion,
     set_event_completion_chunk,
@@ -35,7 +36,7 @@ def traced_method(wrapped, instance, args, kwargs):
         }
         attributes = LLMSpanAttributes(**span_attributes)
         span = tracer.start_span(
-            name=name,
+            name=get_span_name(name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
@@ -76,7 +77,7 @@ async def traced_method(wrapped, instance, args, kwargs):
         }
         attributes = LLMSpanAttributes(**span_attributes)
         span = tracer.start_span(
-            name=name,
+            name=get_span_name(name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
diff --git a/src/langtrace_python_sdk/instrumentation/groq/patch.py b/src/langtrace_python_sdk/instrumentation/groq/patch.py
@@ -29,6 +29,7 @@
     get_llm_request_attributes,
     get_llm_url,
     get_langtrace_attributes,
+    get_span_name,
     set_event_completion,
     set_event_completion_chunk,
     set_usage_attributes,
@@ -107,7 +108,7 @@ def traced_method(wrapped, instance, args, kwargs):
         # with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"],
         #                                   kind=SpanKind.CLIENT) as span:
         span = tracer.start_span(
-            APIS["CHAT_COMPLETION"]["METHOD"],
+            name=get_span_name(APIS["CHAT_COMPLETION"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
@@ -335,7 +336,7 @@ async def traced_method(wrapped, instance, args, kwargs):
         # with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"],
         #                                   kind=SpanKind.CLIENT) as span:
         span = tracer.start_span(
-            APIS["CHAT_COMPLETION"]["METHOD"], kind=SpanKind.CLIENT
+            name=get_span_name(APIS["CHAT_COMPLETION"]["METHOD"]), kind=SpanKind.CLIENT
         )
         for field, value in attributes.model_dump(by_alias=True).items():
             set_span_attribute(span, field, value)
diff --git a/src/langtrace_python_sdk/instrumentation/langchain/patch.py b/src/langtrace_python_sdk/instrumentation/langchain/patch.py
@@ -18,6 +18,7 @@
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.utils.llm import get_span_name
 from opentelemetry import baggage, trace
 from opentelemetry.trace.propagation import set_span_in_context
 from opentelemetry.trace import SpanKind, StatusCode
@@ -57,7 +58,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = FrameworkSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            method_name,
+            name=get_span_name(method_name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/langchain_community/patch.py b/src/langtrace_python_sdk/instrumentation/langchain_community/patch.py
@@ -17,6 +17,7 @@
 import json
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
+from langtrace_python_sdk.utils.llm import get_span_name
 from opentelemetry import baggage, trace
 from opentelemetry.trace.propagation import set_span_in_context
 
@@ -55,7 +56,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = FrameworkSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            method_name,
+            name=get_span_name(method_name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py b/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py
@@ -17,6 +17,7 @@
 import json
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
+from langtrace_python_sdk.utils.llm import get_span_name
 from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind, StatusCode
 from opentelemetry.trace.status import Status
@@ -75,7 +76,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = FrameworkSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            method_name,
+            name=get_span_name(method_name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/langgraph/patch.py b/src/langtrace_python_sdk/instrumentation/langgraph/patch.py
@@ -15,6 +15,7 @@
 """
 
 import json
+from langtrace_python_sdk.utils.llm import get_span_name
 from opentelemetry.trace.propagation import set_span_in_context
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
@@ -52,7 +53,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = FrameworkSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            method_name,
+            name=get_span_name(method_name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/llamaindex/patch.py b/src/langtrace_python_sdk/instrumentation/llamaindex/patch.py
@@ -15,6 +15,7 @@
 """
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
+from langtrace_python_sdk.utils.llm import get_span_name
 from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
@@ -49,7 +50,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = FrameworkSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            method,
+            name=get_span_name(method),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
@@ -94,7 +95,9 @@ async def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            name=get_span_name(method), kind=SpanKind.CLIENT
+        ) as span:
             async for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
                     span.set_attribute(field, value)
diff --git a/src/langtrace_python_sdk/instrumentation/ollama/patch.py b/src/langtrace_python_sdk/instrumentation/ollama/patch.py
@@ -5,6 +5,7 @@
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
+    get_span_name,
     set_event_completion,
     set_event_completion_chunk,
 )
@@ -35,7 +36,7 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = LLMSpanAttributes(**span_attributes)
         with tracer.start_as_current_span(
-            f'ollama.{api["METHOD"]}', kind=SpanKind.CLIENT
+            name=get_span_name(f'ollama.{api["METHOD"]}'), kind=SpanKind.CLIENT
         ) as span:
             _set_input_attributes(span, kwargs, attributes)
 
diff --git a/src/langtrace_python_sdk/instrumentation/openai/patch.py b/src/langtrace_python_sdk/instrumentation/openai/patch.py
@@ -37,6 +37,7 @@
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
+    get_span_name,
     get_tool_calls,
     is_streaming,
     set_event_completion,
@@ -64,7 +65,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_GENERATION"]["METHOD"],
+            name=get_span_name(APIS["IMAGES_GENERATION"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
@@ -127,7 +128,7 @@ async def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_GENERATION"]["METHOD"],
+            name=get_span_name(APIS["IMAGES_GENERATION"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
@@ -192,7 +193,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_EDIT"]["METHOD"],
+            name=APIS["IMAGES_EDIT"]["METHOD"],
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
@@ -282,7 +283,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         span = tracer.start_span(
-            APIS["CHAT_COMPLETION"]["METHOD"],
+            name=get_span_name(APIS["CHAT_COMPLETION"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
@@ -376,7 +377,7 @@ async def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         span = tracer.start_span(
-            APIS["CHAT_COMPLETION"]["METHOD"],
+            name=get_span_name(APIS["CHAT_COMPLETION"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
@@ -455,7 +456,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["EMBEDDINGS_CREATE"]["METHOD"],
+            name=get_span_name(APIS["EMBEDDINGS_CREATE"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
@@ -512,7 +513,7 @@ async def traced_method(wrapped, instance, args, kwargs):
             )
 
         with tracer.start_as_current_span(
-            APIS["EMBEDDINGS_CREATE"]["METHOD"],
+            name=get_span_name(APIS["EMBEDDINGS_CREATE"]["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/pinecone/patch.py b/src/langtrace_python_sdk/instrumentation/pinecone/patch.py
@@ -17,6 +17,7 @@
 import json
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
+from langtrace_python_sdk.utils.llm import get_span_name
 from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
@@ -57,7 +58,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = DatabaseSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            api["METHOD"],
+            name=get_span_name(api["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/qdrant/patch.py b/src/langtrace_python_sdk/instrumentation/qdrant/patch.py
@@ -16,6 +16,7 @@
 
 import json
 from langtrace.trace_attributes import DatabaseSpanAttributes
+from langtrace_python_sdk.utils.llm import get_span_name
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from langtrace_python_sdk.utils import set_span_attribute
 from opentelemetry import baggage, trace
@@ -58,7 +59,7 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = DatabaseSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            api["METHOD"],
+            name=get_span_name(api["METHOD"]),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         ) as span:
diff --git a/src/langtrace_python_sdk/instrumentation/vertexai/patch.py b/src/langtrace_python_sdk/instrumentation/vertexai/patch.py
@@ -8,6 +8,7 @@
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
+    get_span_name,
     set_event_completion,
     set_span_attributes,
     set_usage_attributes,
@@ -39,7 +40,7 @@ def traced_method(wrapped, instance, args, kwargs):
         }
         attributes = LLMSpanAttributes(**span_attributes)
         span = tracer.start_span(
-            name=name,
+            name=get_span_name(name),
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
diff --git a/src/langtrace_python_sdk/instrumentation/weaviate/patch.py b/src/langtrace_python_sdk/instrumentation/weaviate/patch.py
diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py

Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@`
`10`	`10`	`get_langtrace_attributes,`
`11`	`11`	`get_llm_request_attributes,`
`12`	`12`	`get_llm_url,`
	`13`	`+ get_span_name,`
`13`	`14`	`is_streaming,`
`14`	`15`	`set_event_completion,`
`15`	`16`	`set_event_completion_chunk,`
`@@ -35,7 +36,7 @@ def traced_method(wrapped, instance, args, kwargs):`
`35`	`36`	`}`
`36`	`37`	`attributes = LLMSpanAttributes(**span_attributes)`
`37`	`38`	`span = tracer.start_span(`
`38`		`- name=name,`
	`39`	`+ name=get_span_name(name),`
`39`	`40`	`kind=SpanKind.CLIENT,`
`40`	`41`	`context=set_span_in_context(trace.get_current_span()),`
`41`	`42`	`)`
`@@ -76,7 +77,7 @@ async def traced_method(wrapped, instance, args, kwargs):`
`76`	`77`	`}`
`77`	`78`	`attributes = LLMSpanAttributes(**span_attributes)`
`78`	`79`	`span = tracer.start_span(`
`79`		`- name=name,`
	`80`	`+ name=get_span_name(name),`
`80`	`81`	`kind=SpanKind.CLIENT,`
`81`	`82`	`context=set_span_in_context(trace.get_current_span()),`
`82`	`83`	`)`