fix issues with incorrect typing causing tests to fail

Author: rohit-kadhe

.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,
 }

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

@@ -23,12 +23,12 @@
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
 from typing import Any
-from src.langtrace_python_sdk.instrumentation.anthropic.patch import messages_create
+from langtrace_python_sdk.instrumentation.anthropic.patch import messages_create
 
 logging.basicConfig(level=logging.FATAL)
 
 
-class AnthropicInstrumentation(BaseInstrumentor): # type: ignore[misc]
+class AnthropicInstrumentation(BaseInstrumentor):  # type: ignore[misc]
     """
     The AnthropicInstrumentation class represents the Anthropic instrumentation.
     """
@@ -37,7 +37,7 @@ def instrumentation_dependencies(self) -> Collection[str]:
         return ["anthropic >= 0.19.1"]
 
     def _instrument(self, **kwargs: dict[str, Any]) -> None:
-        tracer_provider: TracerProvider = kwargs.get("tracer_provider") # type: ignore
+        tracer_provider: TracerProvider = kwargs.get("tracer_provider")  # type: ignore
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version("anthropic")
 

src/langtrace_python_sdk/instrumentation/anthropic/patch.py

@@ -27,70 +27,47 @@
     get_llm_url,
     get_span_name,
     set_event_completion,
-    set_event_completion_chunk,
     set_usage_attributes,
-    set_span_attribute
+    set_span_attribute,
 )
 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)
+from langtrace_python_sdk.constants.instrumentation.anthropic import APIS
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace_python_sdk.instrumentation.anthropic.types import (
+    StreamingResult,
+    ResultType,
+    MessagesCreateKwargs,
+    ContentItem,
+    Usage,
+)
 
-        span.set_status(StatusCode.OK)
-        span.end()
 
 def messages_create(version: str, tracer: Tracer) -> Callable[..., Any]:
     """Wrap the `messages_create` method."""
 
-    def traced_method(wrapped: Callable[..., Any], instance: Any, args: List[Any], kwargs: MessagesCreateKwargs) -> Any:
+    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
         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"],
-            **extraAttributes,
+            **extraAttributes,  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -103,7 +80,7 @@ def traced_method(wrapped: Callable[..., Any], instance: Any, args: List[Any], k
         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
@@ -114,77 +91,36 @@ def traced_method(wrapped: Callable[..., Any], instance: Any, args: List[Any], k
             span.end()
             raise
 
-    def handle_streaming_response(result: StreamingResult, span: Span) -> Iterator[str]:
-        """Process and yield streaming response chunks."""
-        result_content: List[str] = []
-        span.add_event(Event.STREAM_START.value)
-        input_tokens: int = 0
-        output_tokens: int = 0
-        try:
-            for chunk in result:
-                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 chunk["message"]["usage"] is not None:
-                    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")
-                        else 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 = [{"role": "assistant", "content": "".join(result_content)}]
-            set_event_completion(span, completion)
-
-            span.set_status(StatusCode.OK)
-            span.end()
-
-    def set_response_attributes(result: Union[ResultType, StreamingResult], span: Span, kwargs: MessagesCreateKwargs) -> Any:
+    def set_response_attributes(
+        result: Union[ResultType, StreamingResult], span: Span
+    ) -> Any:
         if not isinstance(result, Iterator):
-            if result["content"] is not None:
+            if hasattr(result, "content") and result.content is not None:
                 set_span_attribute(
-                    span, SpanAttributes.LLM_RESPONSE_MODEL, result["model"]
+                    span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
                 )
-                content_item = result["content"][0]
-                completion: List[ContentItem] = [
+                content_item = result.content[0]
+                completion = [
                     {
-                        "role": result["role"] or "assistant",
-                        "content": content_item.get("text", ""),
-                        "type": content_item.get("type", ""),
+                        "role": result.role or "assistant",
+                        "content": content_item.text,
+                        "type": content_item.type,
                     }
                 ]
                 set_event_completion(span, completion)
 
-            else:
-                responses: List[ContentItem] = []
-                set_event_completion(span, responses)
-
-            if result["system_fingerprint"] is not None:
+            if (
+                hasattr(result, "system_fingerprint")
+                and result.system_fingerprint is not None
+            ):
                 span.set_attribute(
                     SpanAttributes.LLM_SYSTEM_FINGERPRINT,
-                    result["system_fingerprint"],
+                    result.system_fingerprint,
                 )
             # Get the usage
-            if result["usage"] is not None:
-                usage: Usage = result["usage"]
-                set_usage_attributes(span, dict(usage))
+            if hasattr(result, "usage") and result.usage is not None:
+                usage = result.usage
+                set_usage_attributes(span, vars(usage))
 
             span.set_status(StatusCode.OK)
             span.end()

src/langtrace_python_sdk/instrumentation/anthropic/types.py

@@ -12,37 +12,94 @@
 """
 
 from typing import Dict, List, Optional, Iterator, TypedDict
+
+
 class MessagesCreateKwargs(TypedDict, total=False):
-    system: str
+    system: Optional[str]
     messages: List[Dict[str, str]]
 
-class Usage(TypedDict, total=True):
+
+class Usage:
     input_tokens: int
     output_tokens: int
 
-class Message(TypedDict, total=True):
+    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(TypedDict, total=True):
+
+class Delta:
     text: Optional[str]
 
-class Chunk(TypedDict, total=True):
+    def __init__(
+        self,
+        text: Optional[str],
+    ):
+        self.text = text
+
+
+class Chunk:
     message: Message
     delta: Delta
 
-class ContentItem(TypedDict, total=False):
+    def __init__(
+        self,
+        message: Message,
+        delta: Delta,
+    ):
+        self.message = message
+        self.delta = delta
+
+
+class ContentItem:
     role: str
     content: str
     text: str
     type: str
 
-class ResultType(TypedDict, total=True):
+    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]
+StreamingResult = Iterator[Chunk]