add anthropic typings

Author: rohit-kadhe

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/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 langtrace_python_sdk.instrumentation.anthropic.patch import messages_create
+from typing import Any
+from src.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,51 +14,78 @@
 limitations under the License.
 """
 
-import json
-
-from langtrace.trace_attributes import Event, LLMSpanAttributes
-from langtrace_python_sdk.utils import set_span_attribute, silently_fail
+from typing import Any, Callable, Dict, List, Optional, Iterator, TypedDict, Union
+from langtrace.trace_attributes import Event, SpanAttributes, LLMSpanAttributes
 from langtrace_python_sdk.utils.llm import (
     get_extra_attributes,
     get_langtrace_attributes,
     get_llm_request_attributes,
     get_llm_url,
-    is_streaming,
     set_event_completion,
     set_event_completion_chunk,
     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 langtrace_python_sdk.constants.instrumentation.anthropic import APIS
-from langtrace_python_sdk.constants.instrumentation.common import (
-    SERVICE_PROVIDERS,
-)
+from opentelemetry.trace import Span, Tracer, SpanKind
+from opentelemetry.trace.status import StatusCode
+from src.langtrace_python_sdk.constants.instrumentation.anthropic import APIS
+from src.langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from src.langtrace_python_sdk.instrumentation.anthropic.types import StreamingResult, ResultType, MessagesCreateKwargs, ContentItem, Usage
+
+def handle_streaming_response(result: StreamingResult, span: Span) -> Iterator[str]:
+    result_content: List[str] = []
+    span.add_event(Event.STREAM_START.value)
+    input_tokens: int = 0
+    output_tokens: int = 0
+    try:
+        for chunk in result:
+            if chunk['message']["model"] is not None:
+                span.set_attribute(
+                    SpanAttributes.LLM_RESPONSE_MODEL, chunk["message"]["model"]
+                )
+            content: str = ""
+            if chunk["delta"].get("text") is not None:
+                content = chunk["delta"]["text"] or ""
+            result_content.append(content if len(content) > 0 else "")
+
+            if chunk["message"]["usage"] is not None:
+                usage = chunk["message"]["usage"]
+                input_tokens += usage.get("input_tokens", 0)
+                output_tokens += usage.get("output_tokens", 0)
+
+            if content:
+                set_event_completion_chunk(span, "".join(content))
+
+            yield content
+    finally:
+        span.add_event(Event.STREAM_END.value)
+        set_usage_attributes(
+            span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
+        )
+        completion: List[Dict[str, str]] = [{"role": "assistant", "content": "".join(result_content)}]
+        set_event_completion(span, completion)
 
+        span.set_status(StatusCode.OK)
+        span.end()
 
-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", []
-            )
-
+            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,
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -77,56 +104,43 @@ def traced_method(wrapped, instance, args, kwargs):
             # 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
 
-    def handle_streaming_response(result, span):
+    def handle_streaming_response(result: StreamingResult, span: Span) -> Iterator[str]:
         """Process and yield streaming response chunks."""
-        result_content = []
+        result_content: List[str] = []
         span.add_event(Event.STREAM_START.value)
-        input_tokens = 0
-        output_tokens = 0
+        input_tokens: int = 0
+        output_tokens: int = 0
         try:
             for chunk in result:
-                if (
-                    hasattr(chunk, "message")
-                    and chunk.message is not None
-                    and hasattr(chunk.message, "model")
-                    and chunk.message.model is not None
-                ):
-                    span.set_attribute(
-                        SpanAttributes.LLM_RESPONSE_MODEL, chunk.message.model
-                    )
-                content = ""
-                if hasattr(chunk, "delta") and chunk.delta is not None:
-                    content = chunk.delta.text if hasattr(chunk.delta, "text") else ""
-                # Assuming content needs to be aggregated before processing
+                span.set_attribute(
+                    SpanAttributes.LLM_RESPONSE_MODEL, chunk["message"]["model"] or ""
+                )
+                content: str = ""
+                if hasattr(chunk, "delta") and chunk["delta"] is not None:
+                    content = chunk["delta"]["text"] or ""
                 result_content.append(content if len(content) > 0 else "")
-
-                if hasattr(chunk, "message") and hasattr(chunk.message, "usage"):
+                if chunk["message"]["usage"] is not None:
                     input_tokens += (
-                        chunk.message.usage.input_tokens
-                        if hasattr(chunk.message.usage, "input_tokens")
+                        chunk["message"]["usage"]["input_tokens"]
+                        if hasattr(chunk["message"]["usage"], "input_tokens")
                         else 0
                     )
                     output_tokens += (
-                        chunk.message.usage.output_tokens
-                        if hasattr(chunk.message.usage, "output_tokens")
+                        chunk["message"]["usage"]["output_tokens"]
+                        if hasattr(chunk["message"]["usage"], "output_tokens")
                         else 0
                     )
 
-                # Assuming span.add_event is part of a larger logging or event system
-                # Add event for each chunk of content
                 if content:
                     set_event_completion_chunk(span, "".join(content))
 
-                # Assuming this is part of a generator, yield chunk or aggregated content
                 yield content
         finally:
-
-            # Finalize span after processing all chunks
             span.add_event(Event.STREAM_END.value)
             set_usage_attributes(
                 span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
@@ -137,36 +151,34 @@ def handle_streaming_response(result, span):
             span.set_status(StatusCode.OK)
             span.end()
 
-    def set_response_attributes(result, span, kwargs):
-        if not is_streaming(kwargs):
-            if hasattr(result, "content") and result.content is not None:
+    def set_response_attributes(result: Union[ResultType, StreamingResult], span: Span, kwargs: MessagesCreateKwargs) -> Any:
+        if not isinstance(result, Iterator):
+            if result["content"] is not None:
                 set_span_attribute(
-                    span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
+                    span, SpanAttributes.LLM_RESPONSE_MODEL, result["model"]
                 )
-                completion = [
+                content_item = result["content"][0]
+                completion: List[ContentItem] = [
                     {
-                        "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.get("text", ""),
+                        "type": content_item.get("type", ""),
                     }
                 ]
                 set_event_completion(span, completion)
 
             else:
-                responses = []
+                responses: List[ContentItem] = []
                 set_event_completion(span, responses)
 
-            if (
-                hasattr(result, "system_fingerprint")
-                and result.system_fingerprint is not None
-            ):
+            if result["system_fingerprint"] is not None:
                 span.set_attribute(
                     SpanAttributes.LLM_SYSTEM_FINGERPRINT,
-                    result.system_fingerprint,
+                    result["system_fingerprint"],
                 )
             # Get the usage
-            if hasattr(result, "usage") and result.usage is not None:
-                usage = result.usage
+            if result["usage"] is not None:
+                usage: Usage = result["usage"]
                 set_usage_attributes(span, dict(usage))
 
             span.set_status(StatusCode.OK)

src/langtrace_python_sdk/instrumentation/anthropic/types.py

@@ -0,0 +1,35 @@
+from typing import Dict, List, Optional, Iterator, TypedDict
+class MessagesCreateKwargs(TypedDict, total=False):
+    system: str
+    messages: List[Dict[str, str]]
+
+class Usage(TypedDict, total=True):
+    input_tokens: int
+    output_tokens: int
+
+class Message(TypedDict, total=True):
+    model: Optional[str]
+    usage: Optional[Usage]
+
+class Delta(TypedDict, total=True):
+    text: Optional[str]
+
+class Chunk(TypedDict, total=True):
+    message: Message
+    delta: Delta
+
+class ContentItem(TypedDict, total=False):
+    role: str
+    content: str
+    text: str
+    type: str
+
+class ResultType(TypedDict, total=True):
+    model: Optional[str]
+    role: Optional[str]
+    content: List[ContentItem]
+    system_fingerprint: Optional[str]
+    usage: Optional[Usage]
+
+# The result would be an iterator that yields these Chunk objects
+StreamingResult = Iterator[Chunk]

src/langtrace_python_sdk/utils/llm.py

@@ -14,6 +14,7 @@
 limitations under the License.
 """
 
+from typing import Any, Dict
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 from langtrace_python_sdk.utils import set_span_attribute
 from openai import NOT_GIVEN
@@ -136,9 +137,9 @@ def get_llm_request_attributes(kwargs, prompts=None, model=None, operation_name=
     }
 
 
-def get_extra_attributes():
+def get_extra_attributes() -> Dict[str, Any]:
     extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
-    return extra_attributes or {}
+    return extra_attributes.__dict__ or {}
 
 
 def get_llm_url(instance):