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

langtrace_python_sdk/instrumentation/openai/patch.py

@@ -1,21 +1,5 @@
-"""
-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.
-"""
-
 import json
-
+from typing import Any, Dict, List, Optional, Callable, Awaitable, Union
 from langtrace.trace_attributes import (
     LLMSpanAttributes,
     SpanAttributes,
@@ -23,7 +7,7 @@ from langtrace.trace_attributes import (
 from langtrace_python_sdk.utils import set_span_attribute
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from opentelemetry import trace
-from opentelemetry.trace import SpanKind
+from opentelemetry.trace import SpanKind, Tracer, Span
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants.instrumentation.common import (
@@ -46,20 +30,31 @@ from langtrace_python_sdk.utils.llm import (
 )
 from langtrace_python_sdk.types import NOT_GIVEN
 
+from langtrace_python_sdk.instrumentation.openai.types import (
+    ImagesGenerateKwargs,
+    ChatCompletionsCreateKwargs,
+    EmbeddingsCreateKwargs,
+    ImagesEditKwargs,
+    ResultType,
+    ContentItem,
+)
+
 
-def images_generate(original_method, version, tracer):
+def images_generate(version: str, tracer: Tracer) -> Callable:
     """
     Wrap the `generate` method of the `Images` class to trace it.
     """
 
-    def traced_method(wrapped, instance, args, kwargs):
+    def traced_method(
+        wrapped: Callable, instance: Any, args: List[Any], kwargs: ImagesGenerateKwargs
+    ) -> Any:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
         span_attributes = {
             **get_langtrace_attributes(version, service_provider, vendor_type="llm"),
             **get_llm_request_attributes(kwargs, operation_name="images_generate"),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["IMAGES_GENERATION"]["ENDPOINT"],
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -74,21 +69,17 @@ def images_generate(original_method, version, tracer):
                 # Attempt to call the original method
                 result = wrapped(*args, **kwargs)
                 if not is_streaming(kwargs):
-                    data = (
+                    data: Optional[ContentItem] = (
                         result.data[0]
                         if hasattr(result, "data") and len(result.data) > 0
-                        else {}
+                        else None
                     )
                     response = [
                         {
                             "role": "assistant",
                             "content": {
-                                "url": data.url if hasattr(data, "url") else "",
-                                "revised_prompt": (
-                                    data.revised_prompt
-                                    if hasattr(data, "revised_prompt")
-                                    else ""
-                                ),
+                                "url": getattr(data, "url", ""),
+                                "revised_prompt": getattr(data, "revised_prompt", ""),
                             },
                         }
                     ]
@@ -109,12 +100,14 @@ def images_generate(original_method, version, tracer):
     return traced_method
 
 
-def async_images_generate(original_method, version, tracer):
+def async_images_generate(version: str, tracer: Tracer) -> Callable:
     """
     Wrap the `generate` method of the `Images` class to trace it.
     """
 
-    async def traced_method(wrapped, instance, args, kwargs):
+    async def traced_method(
+        wrapped: Callable, instance: Any, args: List[Any], kwargs: ImagesGenerateKwargs
+    ) -> Awaitable[Any]:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
 
         span_attributes = {
@@ -122,7 +115,7 @@ def async_images_generate(original_method, version, tracer):
             **get_llm_request_attributes(kwargs, operation_name="images_generate"),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["IMAGES_GENERATION"]["ENDPOINT"],
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -137,21 +130,17 @@ def async_images_generate(original_method, version, tracer):
                 # Attempt to call the original method
                 result = await wrapped(*args, **kwargs)
                 if not is_streaming(kwargs):
-                    data = (
+                    data: Optional[ContentItem] = (
                         result.data[0]
                         if hasattr(result, "data") and len(result.data) > 0
-                        else {}
+                        else None
                     )
                     response = [
                         {
                             "role": "assistant",
                             "content": {
-                                "url": data.url if hasattr(data, "url") else "",
-                                "revised_prompt": (
-                                    data.revised_prompt
-                                    if hasattr(data, "revised_prompt")
-                                    else ""
-                                ),
+                                "url": getattr(data, "url", ""),
+                                "revised_prompt": getattr(data, "revised_prompt", ""),
                             },
                         }
                     ]
@@ -172,12 +161,14 @@ def async_images_generate(original_method, version, tracer):
     return traced_method
 
 
-def images_edit(original_method, version, tracer):
+def images_edit(version: str, tracer: Tracer) -> Callable:
     """
     Wrap the `edit` method of the `Images` class to trace it.
     """
 
-    def traced_method(wrapped, instance, args, kwargs):
+    def traced_method(
+        wrapped: Callable, instance: Any, args: List[Any], kwargs: ImagesEditKwargs
+    ) -> Any:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
 
         span_attributes = {
@@ -187,7 +178,7 @@ def images_edit(original_method, version, tracer):
             SpanAttributes.LLM_PATH: APIS["IMAGES_EDIT"]["ENDPOINT"],
             SpanAttributes.LLM_RESPONSE_FORMAT: kwargs.get("response_format"),
             SpanAttributes.LLM_IMAGE_SIZE: kwargs.get("size"),
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -233,10 +224,15 @@ def images_edit(original_method, version, tracer):
     return traced_method
 
 
-def chat_completions_create(original_method, version, tracer):
+def chat_completions_create(version: str, tracer: Tracer) -> Callable:
     """Wrap the `create` method of the `ChatCompletion` class to trace it."""
 
-    def traced_method(wrapped, instance, args, kwargs):
+    def traced_method(
+        wrapped: Callable,
+        instance: Any,
+        args: List[Any],
+        kwargs: ChatCompletionsCreateKwargs,
+    ) -> Any:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
         if "perplexity" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["PPLX"]
@@ -251,21 +247,13 @@ def chat_completions_create(original_method, version, tracer):
                 tool_calls = []
                 for tool_call in tools:
                     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)
@@ -277,7 +265,7 @@ def chat_completions_create(original_method, version, tracer):
             **get_llm_request_attributes(kwargs, prompts=llm_prompts),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["CHAT_COMPLETION"]["ENDPOINT"],
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -297,12 +285,9 @@ def chat_completions_create(original_method, version, tracer):
                     prompt_tokens += calculate_prompt_tokens(
                         json.dumps(str(message)), kwargs.get("model")
                     )
-
-                if (
-                    kwargs.get("functions") is not None
-                    and kwargs.get("functions") != NOT_GIVEN
-                ):
-                    for function in kwargs.get("functions"):
+                functions = kwargs.get("functions")
+                if functions is not None and functions != NOT_GIVEN:
+                    for function in functions:
                         prompt_tokens += calculate_prompt_tokens(
                             json.dumps(function), kwargs.get("model")
                         )
@@ -315,7 +300,7 @@ def chat_completions_create(original_method, version, tracer):
                     tool_calls=kwargs.get("tools") is not None,
                 )
             else:
-                _set_response_attributes(span, kwargs, result)
+                _set_response_attributes(span, result)
                 span.set_status(StatusCode.OK)
                 span.end()
                 return result
@@ -329,10 +314,15 @@ def chat_completions_create(original_method, version, tracer):
     return traced_method
 
 
-def async_chat_completions_create(original_method, version, tracer):
+def async_chat_completions_create(version: str, tracer: Tracer) -> Callable:
     """Wrap the `create` method of the `ChatCompletion` class to trace it."""
 
-    async def traced_method(wrapped, instance, args, kwargs):
+    async def traced_method(
+        wrapped: Callable,
+        instance: Any,
+        args: List[Any],
+        kwargs: ChatCompletionsCreateKwargs,
+    ) -> Awaitable[Any]:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
         if "perplexity" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["PPLX"]
@@ -345,21 +335,13 @@ def async_chat_completions_create(original_method, version, tracer):
                 tool_calls = []
                 for tool_call in tools:
                     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(json.dumps(tool_call_dict))
                 llm_prompts.append(tool_calls)
@@ -371,7 +353,7 @@ def async_chat_completions_create(original_method, version, tracer):
             **get_llm_request_attributes(kwargs, prompts=llm_prompts),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["CHAT_COMPLETION"]["ENDPOINT"],
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -392,11 +374,9 @@ def async_chat_completions_create(original_method, version, tracer):
                         json.dumps((str(message))), kwargs.get("model")
                     )
 
-                if (
-                    kwargs.get("functions") is not None
-                    and kwargs.get("functions") != NOT_GIVEN
-                ):
-                    for function in kwargs.get("functions"):
+                functions = kwargs.get("functions")
+                if functions is not None and functions != NOT_GIVEN:
+                    for function in functions:
                         prompt_tokens += calculate_prompt_tokens(
                             json.dumps(function), kwargs.get("model")
                         )
@@ -407,9 +387,9 @@ def async_chat_completions_create(original_method, version, tracer):
                     prompt_tokens,
                     function_call=kwargs.get("functions") is not None,
                     tool_calls=kwargs.get("tools") is not None,
-                )
+                )  # type: ignore
             else:
-                _set_response_attributes(span, kwargs, result)
+                _set_response_attributes(span, result)
                 span.set_status(StatusCode.OK)
                 span.end()
                 return result
@@ -423,12 +403,17 @@ def async_chat_completions_create(original_method, version, tracer):
     return traced_method
 
 
-def embeddings_create(original_method, version, tracer):
+def embeddings_create(version: str, tracer: Tracer) -> Callable:
     """
     Wrap the `create` method of the `Embeddings` class to trace it.
     """
 
-    def traced_method(wrapped, instance, args, kwargs):
+    def traced_method(
+        wrapped: Callable,
+        instance: Any,
+        args: List[Any],
+        kwargs: EmbeddingsCreateKwargs,
+    ) -> Any:
         service_provider = SERVICE_PROVIDERS["OPENAI"]
 
         span_attributes = {
@@ -437,7 +422,7 @@ def embeddings_create(original_method, version, tracer):
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["EMBEDDINGS_CREATE"]["ENDPOINT"],
             SpanAttributes.LLM_REQUEST_DIMENSIONS: kwargs.get("dimensions"),
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         encoding_format = kwargs.get("encoding_format")
@@ -480,12 +465,17 @@ def embeddings_create(original_method, version, tracer):
     return traced_method
 
 
-def async_embeddings_create(original_method, version, tracer):
+def async_embeddings_create(version: str, tracer: Tracer) -> Callable:
     """
     Wrap the `create` method of the `Embeddings` class to trace it.
     """
 
-    async def traced_method(wrapped, instance, args, kwargs):
+    async def traced_method(
+        wrapped: Callable,
+        instance: Any,
+        args: List[Any],
+        kwargs: EmbeddingsCreateKwargs,
+    ) -> Awaitable[Any]:
 
         service_provider = SERVICE_PROVIDERS["OPENAI"]
 
@@ -494,7 +484,7 @@ def async_embeddings_create(original_method, version, tracer):
             **get_llm_request_attributes(kwargs, operation_name="embed"),
             SpanAttributes.LLM_PATH: APIS["EMBEDDINGS_CREATE"]["ENDPOINT"],
             SpanAttributes.LLM_REQUEST_DIMENSIONS: kwargs.get("dimensions"),
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -537,7 +527,7 @@ def async_embeddings_create(original_method, version, tracer):
     return traced_method
 
 
-def extract_content(choice):
+def extract_content(choice: Any) -> Union[str, List[Dict[str, Any]], Dict[str, Any]]:
     # Check if choice.message exists and has a content attribute
     if (
         hasattr(choice, "message")
@@ -582,13 +572,15 @@ def extract_content(choice):
 
 
 @silently_fail
-def _set_input_attributes(span, kwargs, attributes):
+def _set_input_attributes(
+    span: Span, kwargs: ChatCompletionsCreateKwargs, attributes: LLMSpanAttributes
+) -> None:
     tools = []
     for field, value in attributes.model_dump(by_alias=True).items():
         set_span_attribute(span, field, value)
-
-    if kwargs.get("functions") is not None and kwargs.get("functions") != NOT_GIVEN:
-        for function in kwargs.get("functions"):
+    functions = kwargs.get("functions")
+    if functions is not None and functions != NOT_GIVEN:
+        for function in functions:
             tools.append(json.dumps({"type": "function", "function": function}))
 
     if kwargs.get("tools") is not None and kwargs.get("tools") != NOT_GIVEN:
@@ -599,7 +591,7 @@ def _set_input_attributes(span, kwargs, attributes):
 
 
 @silently_fail
-def _set_response_attributes(span, kwargs, result):
+def _set_response_attributes(span: Span, result: ResultType) -> None:
     set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, result.model)
     if hasattr(result, "choices") and result.choices is not None:
         responses = [
@@ -611,8 +603,8 @@ def _set_response_attributes(span, kwargs, result):
                 ),
                 "content": extract_content(choice),
                 **(
-                    {"content_filter_results": choice["content_filter_results"]}
-                    if "content_filter_results" in choice
+                    {"content_filter_results": choice.content_filter_results}
+                    if hasattr(choice, "content_filter_results")
                     else {}
                 ),
             }

langtrace_python_sdk/instrumentation/openai/types.py

@@ -0,0 +1,170 @@
+"""
+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 Any, Dict, List, Union, Optional, TypedDict
+
+
+class ContentItem:
+    url: str
+    revised_prompt: str
+    base64: Optional[str]
+
+    def __init__(
+        self,
+        url: str,
+        revised_prompt: str,
+        base64: Optional[str],
+    ):
+        self.url = url
+        self.revised_prompt = revised_prompt
+        self.base64 = base64
+
+
+class ToolFunction:
+    name: str
+    arguments: str
+
+    def __init__(
+        self,
+        name: str,
+        arguments: str,
+    ):
+        self.name = name
+        self.arguments = arguments
+
+
+class ToolCall:
+    id: str
+    type: str
+    function: ToolFunction
+
+    def __init__(
+        self,
+        id: str,
+        type: str,
+        function: ToolFunction,
+    ):
+        self.id = id
+        self.type = type
+        self.function = function
+
+
+class Message:
+    role: str
+    content: Union[str, List[ContentItem], Dict[str, Any]]
+    tool_calls: Optional[List[ToolCall]]
+
+    def __init__(
+        self,
+        role: str,
+        content: Union[str, List[ContentItem], Dict[str, Any]],
+        content_filter_results: Optional[Any],
+    ):
+        self.role = role
+        self.content = content
+        self.content_filter_results = content_filter_results
+
+
+class Usage:
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+    def __init__(
+        self,
+        prompt_tokens: int,
+        completion_tokens: int,
+        total_tokens: int,
+    ):
+        self.prompt_tokens = prompt_tokens
+        self.completion_tokens = completion_tokens
+        self.total_tokens = total_tokens
+
+
+class Choice:
+    message: Message
+    content_filter_results: Optional[Any]
+
+    def __init__(
+        self,
+        message: Message,
+        content_filter_results: Optional[Any],
+    ):
+        self.message = message
+        self.content_filter_results = content_filter_results
+
+
+class ResultType:
+    model: Optional[str]
+    content: List[ContentItem]
+    system_fingerprint: Optional[str]
+    usage: Optional[Usage]
+    choices: Optional[List[Choice]]
+    response_format: Optional[str]
+    size: Optional[str]
+    encoding_format: Optional[str]
+
+    def __init__(
+        self,
+        model: Optional[str],
+        role: Optional[str],
+        content: List[ContentItem],
+        system_fingerprint: Optional[str],
+        usage: Optional[Usage],
+        functions: Optional[List[ToolCall]],
+        tools: Optional[List[ToolCall]],
+        choices: Optional[List[Choice]],
+        response_format: Optional[str],
+        size: Optional[str],
+        encoding_format: Optional[str],
+    ):
+        self.model = model
+        self.role = role
+        self.content = content
+        self.system_fingerprint = system_fingerprint
+        self.usage = usage
+        self.functions = functions
+        self.tools = tools
+        self.choices = choices
+        self.response_format = response_format
+        self.size = size
+        self.encoding_format = encoding_format
+
+
+class ImagesGenerateKwargs(TypedDict, total=False):
+    operation_name: str
+    model: Optional[str]
+    messages: Optional[List[Message]]
+    functions: Optional[List[ToolCall]]
+    tools: Optional[List[ToolCall]]
+    response_format: Optional[str]
+    size: Optional[str]
+    encoding_format: Optional[str]
+
+
+class ImagesEditKwargs(TypedDict, total=False):
+    response_format: Optional[str]
+    size: Optional[str]
+
+
+class ChatCompletionsCreateKwargs(TypedDict, total=False):
+    model: Optional[str]
+    messages: List[Message]
+    functions: Optional[List[ToolCall]]
+    tools: Optional[List[ToolCall]]
+
+
+class EmbeddingsCreateKwargs(TypedDict, total=False):
+    dimensions: Optional[str]
+    input: Union[str, List[str], None]
+    encoding_format: Optional[Union[List[str], str]]

langtrace_python_sdk/instrumentation/vertexai/patch.py

@@ -102,12 +102,9 @@ def is_streaming_response(response):
 
 
 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("publishers/google/models/", "")
-    return llm_model
+        return instance._model_name.replace("models/", "")
+    return getattr(instance, "_model_id", "unknown")
 
 
 def serialize_prompts(args, kwargs):

langtrace_python_sdk/instrumentation/weaviate/patch.py

@@ -96,19 +96,9 @@ def get_response_object_attributes(response_object):
     response_attributes = {
         **response_object.properties,
         "uuid": str(response_object.uuid) if hasattr(response_object, "uuid") else None,
-        "collection": (
-            response_object.collection
-            if hasattr(response_object, "collection")
-            else None
-        ),
-        "vector": (
-            response_object.vector if hasattr(response_object, "vector") else None
-        ),
-        "references": (
-            response_object.references
-            if hasattr(response_object, "references")
-            else None
-        ),
+        "collection": getattr(response_object, "collection", None),
+        "vector": getattr(response_object, "vector", None),
+        "references": getattr(response_object, "references", None),
         "metadata": (
             extract_metadata(response_object.metadata)
             if hasattr(response_object, "metadata")