Merge branch 'main' of github.com:Scale3-Labs/langtrace-python-sdk into release

Author: karthikscale3

src/examples/langchain_example/sagemaker.py

@@ -0,0 +1,81 @@
+import json
+from typing import Dict
+
+import boto3
+from dotenv import load_dotenv
+from langchain.chains.question_answering import load_qa_chain
+from langchain_community.llms.sagemaker_endpoint import (LLMContentHandler,
+                                                         SagemakerEndpoint)
+from langchain_core.documents import Document
+from langchain_core.prompts import PromptTemplate
+
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
+
+# Add the path to the root of the project to the sys.path
+
+load_dotenv()
+
+langtrace.init()
+example_doc_1 = """
+Peter and Elizabeth took a taxi to attend the night party in the city. While in the party, Elizabeth collapsed and was rushed to the hospital.
+Since she was diagnosed with a brain injury, the doctor told Peter to stay besides her until she gets well.
+Therefore, Peter stayed with her at the hospital for 3 days without leaving.
+"""
+
+docs = [
+    Document(
+        page_content=example_doc_1,
+    )
+]
+
+
+query = """How long was Elizabeth hospitalized?"""
+prompt_template = """Use the following pieces of context to answer the question at the end.
+
+{context}
+
+Question: {question}
+Answer:"""
+PROMPT = PromptTemplate(
+    template=prompt_template, input_variables=["context", "question"]
+)
+
+
+client = boto3.client(
+    "sagemaker-runtime",
+    region_name="us-east-1",
+)
+
+
+class ContentHandler(LLMContentHandler):
+    content_type = "application/json"
+    accepts = "application/json"
+
+    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
+        input_str = json.dumps({"inputs": prompt, "parameters": model_kwargs})
+        return input_str.encode("utf-8")
+
+    def transform_output(self, output: bytes) -> str:
+        response_json = json.loads(output.read().decode("utf-8"))
+        return response_json["generated_text"]
+
+
+@with_langtrace_root_span("SagemakerEndpoint")
+def main():
+    content_handler = ContentHandler()
+
+    chain = load_qa_chain(
+        llm=SagemakerEndpoint(
+            endpoint_name="jumpstart-dft-meta-textgeneration-l-20240809-083223",
+            client=client,
+            model_kwargs={"temperature": 1e-10},
+            content_handler=content_handler,
+        ),
+        prompt=PROMPT,
+    )
+
+    res = chain({"input_documents": docs, "question": query}, return_only_outputs=True)
+    print(res)
+
+
+main()

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -5,6 +5,7 @@
     "gpt-4-1106-preview": "cl100k_base",
     "gpt-4-1106-vision-preview": "cl100k_base",
     "gpt-4o": "0200k_base",
+    "gpt-4o-mini": "0200k_base",
 }
 
 SERVICE_PROVIDERS = {

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)

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:

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:

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:

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()),
         )

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)

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:

src/langtrace_python_sdk/instrumentation/langchain_community/instrumentation.py

@@ -49,9 +49,8 @@ def patch_module_classes(
         lambda member: inspect.isclass(member) and member.__module__ == module.__name__,
     ):
         # loop through all public methods of the class
-        for method_name, _ in inspect.getmembers(obj, predicate=inspect.isfunction):
-            # Skip private methods
-            if method_name.startswith("_"):
+        for method_name, method in inspect.getmembers(obj, predicate=inspect.isfunction):
+            if method.__qualname__.split('.')[0] != name:
                 continue
             try:
                 method_path = f"{name}.{method_name}"
@@ -82,6 +81,12 @@ def _instrument(self, **kwargs):
 
         # List of modules to patch, with their corresponding patch names
         modules_to_patch = [
+            (
+                "langchain_community.llms.sagemaker_endpoint",
+                "sagemaker_endpoint",
+                True,
+                True,
+            ),
             ("langchain_community.document_loaders.pdf", "load_pdf", True, True),
             ("langchain_community.vectorstores.faiss", "vector_store", False, False),
             ("langchain_community.vectorstores.pgvector", "vector_store", False, False),