Merge pull request #205 from Scale3-Labs/ali/sampling

Ali/sampling
Scale3-Labs · Jun 13, 2024 · b9fd119 · b9fd119
2 parents cbde43b + 7411eab
commit b9fd119
Show file tree

Hide file tree

Showing 22 changed files with 284 additions and 52 deletions.
diff --git a/src/examples/langchain_example/basic.py b/src/examples/langchain_example/basic.py
@@ -8,14 +8,15 @@
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 
 from langtrace_python_sdk import langtrace
-from langtrace_python_sdk.utils.with_root_span import (
-    with_langtrace_root_span,
-    with_additional_attributes,
-)
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from openai import OpenAI
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init()
+langtrace.init(
+    write_spans_to_console=False,
+    disable_tracing_for_functions={"langchain": ["RunnableSequence.invoke"]},
+)
 
 
 def api_call_1():
@@ -29,7 +30,8 @@ def api_call_1():
     output_parser = StrOutputParser()
     chain = prompt | llm | output_parser
     res = chain.invoke({"input": "how can langsmith help with testing?"})
-    print(res)
+    # print(res)
+    return res
 
 
 def api_call_2():
@@ -43,13 +45,26 @@ def api_call_2():
     output_parser = StrOutputParser()
     chain = prompt | llm | output_parser
     res = chain.invoke({"input": "how can langsmith help with testing?"})
-    print(res)
+    # print(res)
+    return res
 
 
 @with_langtrace_root_span()
 def basic_app():
     api_call_1()
-    api_call_2()
+    # api_call_2()
+    client = OpenAI()
+    response = client.chat.completions.create(
+        model="gpt-4",
+        messages=[
+            {"role": "system", "content": "Talk like a pirate"},
+            {"role": "user", "content": "Tell me a story in 3 sentences or less."},
+        ],
+        # stream=True,
+        stream=False,
+    )
+
+    return response
 
 
 @with_langtrace_root_span()

diff --git a/src/examples/llamaindex_example/agent.py b/src/examples/llamaindex_example/agent.py
@@ -12,7 +12,7 @@
 
 nest_asyncio.apply()
 
-langtrace.init(write_spans_to_console=False)
+langtrace.init()
 
 
 def multiply(a: int, b: int) -> int:

diff --git a/src/examples/llamaindex_example/basic.py b/src/examples/llamaindex_example/basic.py
@@ -7,7 +7,11 @@
 _ = load_dotenv(find_dotenv())
 
 
-langtrace.init(write_spans_to_console=False)
+langtrace.init(
+    disable_tracing_for_functions={
+        "open_ai": ["openai.chat.completions.create"],
+    }
+)
 
 
 @with_langtrace_root_span()

diff --git a/src/examples/openai_example/chat_completion.py b/src/examples/openai_example/chat_completion.py
@@ -40,5 +40,5 @@ def chat_completion():
             ]
             result.append(content[0] if len(content) > 0 else "")
 
-    print("".join(result))
+    # print("".join(result))
     return response
diff --git a/src/langtrace_python_sdk/constants/instrumentation/pinecone.py b/src/langtrace_python_sdk/constants/instrumentation/pinecone.py
@@ -12,7 +12,7 @@
         "OPERATION": "query",
     },
     "DELETE": {
-        "METHOD": PineconeMethods.DELETE,
+        "METHOD": PineconeMethods.DELETE.value,
         "ENDPOINT": "/vectors/delete",
         "OPERATION": "delete",
     },

diff --git a/src/langtrace_python_sdk/instrumentation/chroma/patch.py b/src/langtrace_python_sdk/instrumentation/chroma/patch.py
@@ -17,10 +17,10 @@
 from langtrace.trace_attributes import DatabaseSpanAttributes
 from langtrace_python_sdk.utils.llm import set_span_attributes
 from langtrace_python_sdk.utils.silently_fail import silently_fail
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants.instrumentation.chroma import APIS
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -59,7 +59,11 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = DatabaseSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            api["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             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/groq/patch.py b/src/langtrace_python_sdk/instrumentation/groq/patch.py
@@ -17,7 +17,8 @@
 import json
 
 from langtrace.trace_attributes import Event, LLMSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
+from opentelemetry.trace.propagation import set_span_in_context
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
@@ -112,7 +113,9 @@ 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
+            APIS["CHAT_COMPLETION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         )
         for field, value in attributes.model_dump(by_alias=True).items():
             if value is not None:

diff --git a/src/langtrace_python_sdk/instrumentation/langchain/patch.py b/src/langtrace_python_sdk/instrumentation/langchain/patch.py
@@ -17,7 +17,9 @@
 import json
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry import baggage
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from opentelemetry import baggage, trace
+from opentelemetry.trace.propagation import set_span_in_context
 from opentelemetry.trace import SpanKind, StatusCode
 from opentelemetry.trace.status import Status
 
@@ -54,7 +56,11 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            method_name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             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/langchain_community/patch.py b/src/langtrace_python_sdk/instrumentation/langchain_community/patch.py
@@ -17,7 +17,9 @@
 import json
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
+from opentelemetry.trace.propagation import set_span_in_context
+
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
@@ -52,7 +54,11 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            method_name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             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/langchain_core/patch.py b/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py
@@ -17,9 +17,10 @@
 import json
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind, StatusCode
 from opentelemetry.trace.status import Status
+from opentelemetry.trace.propagation import set_span_in_context
 
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -73,7 +74,12 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            method_name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
                     span.set_attribute(field, value)
@@ -147,7 +153,11 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            method_name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             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/langgraph/patch.py b/src/langtrace_python_sdk/instrumentation/langgraph/patch.py
@@ -15,9 +15,10 @@
 """
 
 import json
+from opentelemetry.trace.propagation import set_span_in_context
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
@@ -50,7 +51,11 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            method_name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             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/llamaindex/patch.py b/src/langtrace_python_sdk/instrumentation/llamaindex/patch.py
@@ -15,10 +15,10 @@
 """
 
 from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
@@ -48,7 +48,11 @@ 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(
+            method,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             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/openai/patch.py b/src/langtrace_python_sdk/instrumentation/openai/patch.py
@@ -18,10 +18,10 @@
 
 from importlib_metadata import version as v
 from langtrace.trace_attributes import Event, LLMSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -64,7 +64,9 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_GENERATION"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["IMAGES_GENERATION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
@@ -141,7 +143,9 @@ async def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_GENERATION"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["IMAGES_GENERATION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
             items = attributes.model_dump(by_alias=True).items()
             for field, value in items:
@@ -226,7 +230,9 @@ def traced_method(wrapped, instance, args, kwargs):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_EDIT"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["IMAGES_EDIT"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
@@ -349,7 +355,9 @@ 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
+            APIS["CHAT_COMPLETION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         )
         for field, value in attributes.model_dump(by_alias=True).items():
             if value is not None:
@@ -833,7 +841,9 @@ def traced_method(wrapped, instance, args, kwargs):
             attributes["llm.user"] = kwargs.get("user")
 
         with tracer.start_as_current_span(
-            APIS["EMBEDDINGS_CREATE"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["EMBEDDINGS_CREATE"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
 
             for field, value in attributes.model_dump(by_alias=True).items():
@@ -898,7 +908,9 @@ async def traced_method(wrapped, instance, args, kwargs):
             attributes["llm.user"] = kwargs.get("user")
 
         with tracer.start_as_current_span(
-            APIS["EMBEDDINGS_CREATE"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["EMBEDDINGS_CREATE"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
 
             async for field, value in attributes.model_dump(by_alias=True).items():

diff --git a/src/langtrace_python_sdk/instrumentation/pinecone/patch.py b/src/langtrace_python_sdk/instrumentation/pinecone/patch.py
@@ -17,10 +17,10 @@
 import json
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
@@ -56,7 +56,11 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = DatabaseSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            api["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
 
             if span.is_recording():
                 set_span_attribute(span, "server.address", instance._config.host)