Merge pull request #326 from Scale3-Labs/release-2.3.4

Release 2.3.4: Add Types for OpenAI & Anthropic Instrumentations + Refactoring accessing attributes

Author: alizenhom

.vscode/settings.json

@@ -1,7 +1,11 @@
-
 {
   "[python]": {
     "editor.defaultFormatter": "ms-python.black-formatter",
   },
   "editor.formatOnSave": true,
+  "python.testing.pytestArgs": [
+    "src"
+  ],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true,
 }

mypy.ini

@@ -0,0 +1,16 @@
+[mypy]
+strict = True
+disable_error_code = import-untyped
+disallow_untyped_calls = True        # Disallow function calls without type annotations
+disallow_untyped_defs = True         # Disallow defining functions without type annotations
+disallow_any_explicit = True         # Disallow explicit use of `Any`
+disallow_any_generics = True         # Disallow generic types without specific type parameters
+disallow_incomplete_defs = True      # Disallow defining incomplete function signatures
+no_implicit_optional = True          # Disallow implicitly Optional types
+warn_unused_configs = True           # Warn about unused configurations
+warn_redundant_casts = True          # Warn about unnecessary type casts
+warn_return_any = True               # Warn if a function returns `Any`
+warn_unreachable = True              # Warn about unreachable code
+# Ignore external modules or allow specific imports
+follow_imports = skip
+ignore_missing_imports = True

src/examples/anthropic_example/completion.py

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

src/examples/openai_example/__init__.py

@@ -20,3 +20,4 @@ def run(self):
         chat_completion_example()
         embeddings_create_example()
         function_example()
+        image_edit()

src/langtrace_python_sdk/instrumentation/anthropic/instrumentation.py

@@ -16,38 +16,39 @@
 
 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

src/langtrace_python_sdk/instrumentation/anthropic/patch.py

@@ -14,9 +14,8 @@
 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 @@
     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 traced_method(wrapped, instance, args, kwargs):
         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 set_response_attributes(result, span, kwargs):
             # 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()

src/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]

src/langtrace_python_sdk/instrumentation/cohere/patch.py

@@ -367,10 +367,7 @@ def traced_method(wrapped, instance, args, kwargs):
                 }
                 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),

src/langtrace_python_sdk/instrumentation/crewai/patch.py

@@ -8,7 +8,9 @@
 
 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 traced_method(wrapped, instance, args, kwargs):
                 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 traced_method(wrapped, instance, args, kwargs):
                 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