langtrace-python-sdk 2.3.3__py3-none-any.whl → 2.3.4__py3-none-any.whl

examples/anthropic_example/completion.py

@@ -7,7 +7,7 @@ from langtrace_python_sdk import langtrace, with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init()
+langtrace.init(write_spans_to_console=True)
 
 
 @with_langtrace_root_span("messages_create")

examples/openai_example/__init__.py

@@ -20,3 +20,4 @@ class OpenAIRunner:
         chat_completion_example()
         embeddings_create_example()
         function_example()
+        image_edit()

langtrace_python_sdk/instrumentation/anthropic/instrumentation.py

@@ -16,38 +16,39 @@ limitations under the License.
 
 import importlib.metadata
 import logging
-from typing import Collection
+from typing import Collection, Any
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import TracerProvider
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
-
+from typing import Any
 from langtrace_python_sdk.instrumentation.anthropic.patch import messages_create
 
 logging.basicConfig(level=logging.FATAL)
 
 
-class AnthropicInstrumentation(BaseInstrumentor):
+class AnthropicInstrumentation(BaseInstrumentor):  # type: ignore[misc]
     """
-    The AnthropicInstrumentation class represents the Anthropic instrumentation
+    The AnthropicInstrumentation class represents the Anthropic instrumentation.
     """
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["anthropic >= 0.19.1"]
 
-    def _instrument(self, **kwargs):
-        tracer_provider = kwargs.get("tracer_provider")
+    def _instrument(self, **kwargs: dict[str, Any]) -> None:
+        tracer_provider: TracerProvider = kwargs.get("tracer_provider")  # type: ignore
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version("anthropic")
 
         wrap_function_wrapper(
             "anthropic.resources.messages",
             "Messages.create",
-            messages_create("anthropic.messages.create", version, tracer),
+            messages_create(version, tracer),
         )
 
-    def _instrument_module(self, module_name):
+    def _instrument_module(self, module_name: str) -> None:
         pass
 
-    def _uninstrument(self, **kwargs):
+    def _uninstrument(self, **kwargs: dict[str, Any]) -> None:
         pass

langtrace_python_sdk/instrumentation/anthropic/patch.py

@@ -14,9 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
-
-from langtrace.trace_attributes import Event, LLMSpanAttributes
+from typing import Any, Callable, Dict, List, Optional, Iterator, TypedDict, Union
+from langtrace.trace_attributes import Event, SpanAttributes, LLMSpanAttributes
 from langtrace_python_sdk.utils import set_span_attribute
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 
@@ -27,41 +26,48 @@ from langtrace_python_sdk.utils.llm import (
     get_llm_request_attributes,
     get_llm_url,
     get_span_name,
-    is_streaming,
     set_event_completion,
     set_usage_attributes,
+    set_span_attribute,
 )
-from opentelemetry.trace import SpanKind
-from opentelemetry.trace.status import Status, StatusCode
-from langtrace.trace_attributes import SpanAttributes
-
+from opentelemetry.trace import Span, Tracer, SpanKind
+from opentelemetry.trace.status import StatusCode
 from langtrace_python_sdk.constants.instrumentation.anthropic import APIS
-from langtrace_python_sdk.constants.instrumentation.common import (
-    SERVICE_PROVIDERS,
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace_python_sdk.instrumentation.anthropic.types import (
+    StreamingResult,
+    ResultType,
+    MessagesCreateKwargs,
+    ContentItem,
+    Usage,
 )
 
 
-def messages_create(original_method, version, tracer):
+def messages_create(version: str, tracer: Tracer) -> Callable[..., Any]:
     """Wrap the `messages_create` method."""
 
-    def traced_method(wrapped, instance, args, kwargs):
+    def traced_method(
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: List[Any],
+        kwargs: MessagesCreateKwargs,
+    ) -> Any:
         service_provider = SERVICE_PROVIDERS["ANTHROPIC"]
 
-        # extract system from kwargs and attach as a role to the prompts
-        # we do this to keep it consistent with the openai
+        # Extract system from kwargs and attach as a role to the prompts
         prompts = kwargs.get("messages", [])
         system = kwargs.get("system")
         if system:
             prompts = [{"role": "system", "content": system}] + kwargs.get(
                 "messages", []
             )
-
+        extraAttributes = get_extra_attributes()
         span_attributes = {
             **get_langtrace_attributes(version, service_provider),
             **get_llm_request_attributes(kwargs, prompts=prompts),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["MESSAGES_CREATE"]["ENDPOINT"],
-            **get_extra_attributes(),
+            **extraAttributes,  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -74,37 +80,35 @@ def messages_create(original_method, version, tracer):
         try:
             # Attempt to call the original method
             result = wrapped(*args, **kwargs)
-            return set_response_attributes(result, span, kwargs)
+            return set_response_attributes(result, span)
 
         except Exception as err:
             # Record the exception in the span
             span.record_exception(err)
             # Set the span status to indicate an error
-            span.set_status(Status(StatusCode.ERROR, str(err)))
+            span.set_status(StatusCode.ERROR, str(err))
             # Reraise the exception to ensure it's not swallowed
             span.end()
             raise
 
-    @silently_fail
-    def set_response_attributes(result, span, kwargs):
-        if not is_streaming(kwargs):
+    def set_response_attributes(
+        result: Union[ResultType, StreamingResult], span: Span
+    ) -> Any:
+        if not isinstance(result, Iterator):
             if hasattr(result, "content") and result.content is not None:
                 set_span_attribute(
                     span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
                 )
+                content_item = result.content[0]
                 completion = [
                     {
-                        "role": result.role if result.role else "assistant",
-                        "content": result.content[0].text,
-                        "type": result.content[0].type,
+                        "role": result.role or "assistant",
+                        "content": content_item.text,
+                        "type": content_item.type,
                     }
                 ]
                 set_event_completion(span, completion)
 
-            else:
-                responses = []
-                set_event_completion(span, responses)
-
             if (
                 hasattr(result, "system_fingerprint")
                 and result.system_fingerprint is not None
@@ -116,7 +120,7 @@ def messages_create(original_method, version, tracer):
             # Get the usage
             if hasattr(result, "usage") and result.usage is not None:
                 usage = result.usage
-                set_usage_attributes(span, dict(usage))
+                set_usage_attributes(span, vars(usage))
 
             span.set_status(StatusCode.OK)
             span.end()

langtrace_python_sdk/instrumentation/anthropic/types.py

@@ -0,0 +1,105 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Dict, List, Optional, Iterator, TypedDict
+
+
+class MessagesCreateKwargs(TypedDict, total=False):
+    system: Optional[str]
+    messages: List[Dict[str, str]]
+
+
+class Usage:
+    input_tokens: int
+    output_tokens: int
+
+    def __init__(self, input_tokens: int, output_tokens: int):
+        self.input_tokens = input_tokens
+        self.output_tokens = output_tokens
+
+
+class Message:
+    def __init__(
+        self,
+        id: str,
+        model: Optional[str],
+        usage: Optional[Usage],
+    ):
+        self.id = id
+        self.model = model
+        self.usage = usage
+
+    model: Optional[str]
+    usage: Optional[Usage]
+
+
+class Delta:
+    text: Optional[str]
+
+    def __init__(
+        self,
+        text: Optional[str],
+    ):
+        self.text = text
+
+
+class Chunk:
+    message: Message
+    delta: Delta
+
+    def __init__(
+        self,
+        message: Message,
+        delta: Delta,
+    ):
+        self.message = message
+        self.delta = delta
+
+
+class ContentItem:
+    role: str
+    content: str
+    text: str
+    type: str
+
+    def __init__(self, role: str, content: str, text: str, type: str):
+        self.role = role
+        self.content = content
+        self.text = text
+        self.type = type
+
+
+class ResultType:
+    model: Optional[str]
+    role: Optional[str]
+    content: List[ContentItem]
+    system_fingerprint: Optional[str]
+    usage: Optional[Usage]
+
+    def __init__(
+        self,
+        model: Optional[str],
+        role: Optional[str],
+        content: Optional[List[ContentItem]],
+        system_fingerprint: Optional[str],
+        usage: Optional[Usage],
+    ):
+        self.model = model
+        self.role = role
+        self.content = content if content is not None else []
+        self.system_fingerprint = system_fingerprint
+        self.usage = usage
+
+
+# The result would be an iterator that yields these Chunk objects
+StreamingResult = Iterator[Chunk]

langtrace_python_sdk/instrumentation/cohere/patch.py

@@ -367,10 +367,7 @@ def chat_stream(original_method, version, tracer):
                 }
                 for item in chat_history
             ]
-        if len(history) > 0:
-            prompts = history + prompts
-        if len(system_prompts) > 0:
-            prompts = system_prompts + prompts
+        prompts = system_prompts + history + prompts
 
         span_attributes = {
             **get_langtrace_attributes(version, service_provider),

langtrace_python_sdk/instrumentation/crewai/patch.py

@@ -8,7 +8,9 @@ from opentelemetry.trace.status import Status, StatusCode
 
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 from langtrace_python_sdk.constants.instrumentation.common import (
-    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY, SERVICE_PROVIDERS)
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
 from langtrace_python_sdk.utils import set_span_attribute
 from langtrace_python_sdk.utils.llm import get_span_name, set_span_attributes
 from langtrace_python_sdk.utils.misc import serialize_args, serialize_kwargs
@@ -44,7 +46,9 @@ def patch_memory(operation_name, version, tracer: Tracer):
                 set_span_attributes(span, attributes)
                 result = wrapped(*args, **kwargs)
                 if result is not None and len(result) > 0:
-                    set_span_attribute(span, "crewai.memory.storage.rag_storage.outputs", str(result))
+                    set_span_attribute(
+                        span, "crewai.memory.storage.rag_storage.outputs", str(result)
+                    )
                 if result:
                     span.set_status(Status(StatusCode.OK))
                 span.end()
@@ -87,20 +91,17 @@ def patch_crew(operation_name, version, tracer: Tracer):
                 CrewAISpanAttributes(span=span, instance=instance)
                 result = wrapped(*args, **kwargs)
                 if result:
+                    class_name = instance.__class__.__name__
+                    span.set_attribute(
+                        f"crewai.{class_name.lower()}.result", str(result)
+                    )
                     span.set_status(Status(StatusCode.OK))
-                    if instance.__class__.__name__ == "Crew":
-                        span.set_attribute("crewai.crew.result", str(result))
-                        if hasattr(result, "tasks_output") and result.tasks_output is not None:
-                            span.set_attribute("crewai.crew.tasks_output", str((result.tasks_output)))
-                        if hasattr(result, "token_usage") and result.token_usage is not None:
-                            span.set_attribute("crewai.crew.token_usage", str((result.token_usage)))
-                        if hasattr(result, "usage_metrics") and result.usage_metrics is not None:
-                            span.set_attribute("crewai.crew.usage_metrics", str((result.usage_metrics)))
-                    elif instance.__class__.__name__ == "Agent":
-                        span.set_attribute("crewai.agent.result", str(result))
-                    elif instance.__class__.__name__ == "Task":
-                        span.set_attribute("crewai.task.result", str(result))
-
+                    if class_name == "Crew":
+                        for attr in ["tasks_output", "token_usage", "usage_metrics"]:
+                            if hasattr(result, attr):
+                                span.set_attribute(
+                                    f"crewai.crew.{attr}", str(getattr(result, attr))
+                                )
                 span.end()
                 return result
 

langtrace_python_sdk/instrumentation/gemini/patch.py

@@ -100,12 +100,9 @@ def apatch_gemini(name, version, tracer: Tracer):
 
 
 def get_llm_model(instance):
-    llm_model = "unknown"
-    if hasattr(instance, "_model_id"):
-        llm_model = instance._model_id
     if hasattr(instance, "_model_name"):
-        llm_model = instance._model_name.replace("models/", "")
-    return llm_model
+        return instance._model_name.replace("models/", "")
+    return getattr(instance, "_model_id", "unknown")
 
 
 def serialize_prompts(args, kwargs, instance):

langtrace_python_sdk/instrumentation/groq/patch.py

@@ -158,7 +158,7 @@ def chat_completions_create(original_method, version, tracer):
                     usage = result.usage
                     set_usage_attributes(span, dict(usage))
 
-                span.set_status(StatusCode.OK)
+                span.set_status(Status(StatusCode.OK))
                 span.end()
                 return result
             else:
@@ -257,7 +257,7 @@ def chat_completions_create(original_method, version, tracer):
                 span, [{"role": "assistant", "content": "".join(result_content)}]
             )
 
-            span.set_status(StatusCode.OK)
+            span.set_status(Status(StatusCode.OK))
             span.end()
 
     # return the wrapped method
@@ -282,21 +282,13 @@ def async_chat_completions_create(original_method, version, tracer):
                 tool_calls = []
                 for tool_call in item.tool_calls:
                     tool_call_dict = {
-                        "id": tool_call.id if hasattr(tool_call, "id") else "",
-                        "type": tool_call.type if hasattr(tool_call, "type") else "",
+                        "id": getattr(tool_call, "id", ""),
+                        "type": getattr(tool_call, "type", ""),
                     }
                     if hasattr(tool_call, "function"):
                         tool_call_dict["function"] = {
-                            "name": (
-                                tool_call.function.name
-                                if hasattr(tool_call.function, "name")
-                                else ""
-                            ),
-                            "arguments": (
-                                tool_call.function.arguments
-                                if hasattr(tool_call.function, "arguments")
-                                else ""
-                            ),
+                            "name": getattr(tool_call.function, "name", ""),
+                            "arguments": getattr(tool_call.function, "arguments", ""),
                         }
                     tool_calls.append(tool_call_dict)
                 llm_prompts.append(tool_calls)
@@ -459,11 +451,7 @@ def async_chat_completions_create(original_method, version, tracer):
                                             tool_call.function.arguments
                                         )
                                         completion_tokens += token_counts
-                                        content = content + [
-                                            tool_call.function.arguments
-                                        ]
-                                    else:
-                                        content = content + []
+                                        content += [tool_call.function.arguments]
                 else:
                     content = []
                 span.add_event(

langtrace_python_sdk/instrumentation/openai/instrumentation.py

@@ -1,12 +1,9 @@
 """
 Copyright (c) 2024 Scale3 Labs
-
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
-
     http://www.apache.org/licenses/LICENSE-2.0
-
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -14,12 +11,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from typing import Collection, Optional, Any
 import importlib.metadata
 import logging
-from typing import Collection
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
-from opentelemetry.trace import get_tracer
+from opentelemetry.trace import get_tracer, TracerProvider
 from wrapt import wrap_function_wrapper
 
 from langtrace_python_sdk.instrumentation.openai.patch import (
@@ -35,59 +32,57 @@ from langtrace_python_sdk.instrumentation.openai.patch import (
 logging.basicConfig(level=logging.FATAL)
 
 
-class OpenAIInstrumentation(BaseInstrumentor):
+class OpenAIInstrumentation(BaseInstrumentor): # type: ignore
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["openai >= 0.27.0", "trace-attributes >= 4.0.5"]
 
-    def _instrument(self, **kwargs):
-        tracer_provider = kwargs.get("tracer_provider")
+    def _instrument(self, **kwargs: Any) -> None:
+        tracer_provider: Optional[TracerProvider] = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, "", tracer_provider)
-        version = importlib.metadata.version("openai")
+        version: str = importlib.metadata.version("openai")
 
         wrap_function_wrapper(
             "openai.resources.chat.completions",
             "Completions.create",
-            chat_completions_create("openai.chat.completions.create", version, tracer),
+            chat_completions_create(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.chat.completions",
             "AsyncCompletions.create",
-            async_chat_completions_create(
-                "openai.chat.completions.create_stream", version, tracer
-            ),
+            async_chat_completions_create(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.images",
             "Images.generate",
-            images_generate("openai.images.generate", version, tracer),
+            images_generate(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.images",
             "AsyncImages.generate",
-            async_images_generate("openai.images.generate", version, tracer),
+            async_images_generate(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.images",
             "Images.edit",
-            images_edit("openai.images.edit", version, tracer),
+            images_edit(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.embeddings",
             "Embeddings.create",
-            embeddings_create("openai.embeddings.create", version, tracer),
+            embeddings_create(version, tracer),
         )
 
         wrap_function_wrapper(
             "openai.resources.embeddings",
             "AsyncEmbeddings.create",
-            async_embeddings_create("openai.embeddings.create", version, tracer),
+            async_embeddings_create(version, tracer),
         )
 
-    def _uninstrument(self, **kwargs):
+    def _uninstrument(self, **kwargs: Any) -> None:
         pass