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

langtrace_python_sdk/instrumentation/crewai/instrumentation.py

@@ -39,6 +39,21 @@ class CrewAIInstrumentation(BaseInstrumentor):
                 "Crew.kickoff",
                 patch_crew("Crew.kickoff", version, tracer),
             )
+            _W(
+                "crewai.crew",
+                "Crew.kickoff_for_each",
+                patch_crew("Crew.kickoff_for_each", version, tracer),
+            )
+            _W(
+                "crewai.crew",
+                "Crew.kickoff_async",
+                patch_crew("Crew.kickoff_async", version, tracer),
+            )
+            _W(
+                "crewai.crew",
+                "Crew.kickoff_for_each_async",
+                patch_crew("Crew.kickoff_for_each_async", version, tracer),
+            )
             _W(
                 "crewai.agent",
                 "Agent.execute_task",

langtrace_python_sdk/instrumentation/crewai/patch.py

@@ -1,16 +1,18 @@
 import json
+
 from importlib_metadata import version as v
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry import baggage
+from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace.status import Status, StatusCode
+
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
-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.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
 )
-from opentelemetry import baggage
-from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry.trace import SpanKind, Span, Tracer
-from opentelemetry.trace.status import Status, StatusCode
+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"):
-                            span.set_attribute("crewai.crew.tasks_output", str((result.tasks_output)))
-                        if hasattr(result, "token_usage"):
-                            span.set_attribute("crewai.crew.token_usage", str((result.token_usage)))
-                        if hasattr(result, "usage_metrics"):
-                            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
 
@@ -137,22 +138,27 @@ class CrewAISpanAttributes:
             self.set_crew_attributes()
             for key, value in self.crew.items():
                 key = f"crewai.crew.{key}"
-                set_span_attribute(self.span, key, value)
+                if value is not None:
+                    set_span_attribute(self.span, key, value)
 
         elif instance_name == "Agent":
             agent = self.set_agent_attributes()
             for key, value in agent.items():
                 key = f"crewai.agent.{key}"
-                set_span_attribute(self.span, key, value)
+                if value is not None:
+                    set_span_attribute(self.span, key, value)
 
         elif instance_name == "Task":
             task = self.set_task_attributes()
             for key, value in task.items():
                 key = f"crewai.task.{key}"
-                set_span_attribute(self.span, key, value)
+                if value is not None:
+                    set_span_attribute(self.span, key, value)
 
     def set_crew_attributes(self):
         for key, value in self.instance.__dict__.items():
+            if value is None:
+                continue
             if key == "tasks":
                 self._parse_tasks(value)
             elif key == "agents":
@@ -163,6 +169,8 @@ class CrewAISpanAttributes:
     def set_agent_attributes(self):
         agent = {}
         for key, value in self.instance.__dict__.items():
+            if key == "tools":
+                value = self._parse_tools(value)
             if value is None:
                 continue
             agent[key] = str(value)
@@ -174,8 +182,10 @@ class CrewAISpanAttributes:
         for key, value in self.instance.__dict__.items():
             if value is None:
                 continue
-
-            if key == "agent":
+            if key == "tools":
+                value = self._parse_tools(value)
+                task[key] = value
+            elif key == "agent":
                 task[key] = value.role
             else:
                 task[key] = str(value)
@@ -218,3 +228,15 @@ class CrewAISpanAttributes:
                     "output_file": task.output_file,
                 }
             )
+
+    def _parse_tools(self, tools):
+        result = []
+        for tool in tools:
+            res = {}
+            if hasattr(tool, "name") and tool.name is not None:
+                res["name"] = tool.name
+            if hasattr(tool, "description") and tool.description is not None:
+                res["description"] = tool.description
+            if res:
+                result.append(res)
+        return json.dumps(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