openlit 1.34.7__py3-none-any.whl → 1.34.10__py3-none-any.whl

openlit/instrumentation/ollama/ollama.py

@@ -2,7 +2,6 @@
 Module for monitoring Ollama API calls.
 """
 
-import logging
 import time
 from opentelemetry.trace import SpanKind
 from openlit.__helpers import (
@@ -17,12 +16,10 @@ from openlit.instrumentation.ollama.utils import (
 )
 from openlit.semcov import SemanticConvention
 
-logger = logging.getLogger(__name__)
-
 def chat(version, environment, application_name,
-            tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
+            tracer, pricing_info, capture_message_content, metrics, disable_metrics):
     """
-    Generates a telemetry wrapper for GenAI function call
+    Generates a telemetry wrapper for Ollama chat function call
     """
 
     class TracedSyncStream:
@@ -38,7 +35,7 @@ def chat(version, environment, application_name,
                 kwargs,
                 server_address,
                 server_port,
-                **args,
+                args,
             ):
             self.__wrapped__ = wrapped
             self._span = span
@@ -48,11 +45,11 @@ def chat(version, environment, application_name,
             self._tool_calls = []
             self._input_tokens = 0
             self._output_tokens = 0
-            self._response_role = ''
+            self._response_role = ""
             self._span_name = span_name
             self._args = args
             self._kwargs = kwargs
-            self._start_time = time.time()
+            self._start_time = time.monotonic()
             self._end_time = None
             self._timestamps = []
             self._ttft = 0
@@ -81,56 +78,101 @@ def chat(version, environment, application_name,
                 return chunk
             except StopIteration:
                 try:
-                    with tracer.start_as_current_span(self._span_name, kind= SpanKind.CLIENT) as self._span:
+                    with tracer.start_as_current_span(self._span_name, kind=SpanKind.CLIENT) as self._span:
                         process_streaming_chat_response(
                             self,
                             pricing_info=pricing_info,
                             environment=environment,
                             application_name=application_name,
                             metrics=metrics,
-                            event_provider=event_provider,
                             capture_message_content=capture_message_content,
                             disable_metrics=disable_metrics,
                             version=version
                         )
                 except Exception as e:
                     handle_exception(self._span, e)
-                    logger.error("Error in trace creation: %s", e)
+
                 raise
 
     def wrapper(wrapped, instance, args, kwargs):
         """
-        Wraps the GenAI function call.
+        Wraps the Ollama chat function call.
         """
 
         streaming = kwargs.get("stream", False)
 
         server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 11434)
-        json_body = kwargs.get("json", {}) or {}
-        request_model = json_body.get("model") or kwargs.get("model")
+        request_model = kwargs.get("model")
 
         span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
 
-        # pylint: disable=no-else-return
         if streaming:
             awaited_wrapped = wrapped(*args, **kwargs)
             span = tracer.start_span(span_name, kind=SpanKind.CLIENT)
-            return TracedSyncStream(awaited_wrapped, span, span_name, kwargs, server_address, server_port)
+            return TracedSyncStream(awaited_wrapped, span, span_name, kwargs, server_address, server_port, args)
 
         else:
-            with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
-                start_time = time.time()
+            with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
+                start_time = time.monotonic()
+
+                try:
+                    response = wrapped(*args, **kwargs)
+
+                    response = process_chat_response(
+                        response=response,
+                        gen_ai_endpoint="ollama.chat",
+                        pricing_info=pricing_info,
+                        server_port=server_port,
+                        server_address=server_address,
+                        environment=environment,
+                        application_name=application_name,
+                        metrics=metrics,
+                        start_time=start_time,
+                        span=span,
+                        capture_message_content=capture_message_content,
+                        disable_metrics=disable_metrics,
+                        version=version,
+                        **kwargs
+                    )
+
+                except Exception as e:
+                    handle_exception(span, e)
+
+            return response
+
+    return wrapper
+
+def embeddings(version, environment, application_name,
+              tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+    """
+    Generates a telemetry wrapper for Ollama embeddings function call
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        """
+        Wraps the Ollama embeddings function call.
+        """
+
+        server_address, server_port = set_server_address_and_port(instance, "127.0.0.1", 11434)
+        request_model = kwargs.get("model")
+
+        span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING} {request_model}"
+
+        with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
+            start_time = time.monotonic()
+
+            try:
                 response = wrapped(*args, **kwargs)
-                response = process_chat_response(
+
+                response = process_embedding_response(
                     response=response,
-                    request_model=request_model,
+                    gen_ai_endpoint="ollama.embeddings",
                     pricing_info=pricing_info,
                     server_port=server_port,
                     server_address=server_address,
                     environment=environment,
                     application_name=application_name,
                     metrics=metrics,
-                    event_provider=event_provider,
                     start_time=start_time,
                     span=span,
                     capture_message_content=capture_message_content,
@@ -139,47 +181,8 @@ def chat(version, environment, application_name,
                     **kwargs
                 )
 
-            return response
-
-    return wrapper
-
-def embeddings(version, environment, application_name,
-              tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
-    """
-    Generates a telemetry wrapper for GenAI function call
-    """
-
-    def wrapper(wrapped, instance, args, kwargs):
-        """
-        Wraps the GenAI function call.
-        """
-
-        server_address, server_port = set_server_address_and_port(instance, '127.0.0.1', 11434)
-        json_body = kwargs.get('json', {}) or {}
-        request_model = json_body.get('model') or kwargs.get('model')
-
-        span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING} {request_model}'
-
-        with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
-            start_time = time.time()
-            response = wrapped(*args, **kwargs)
-            response = process_embedding_response(
-                response=response,
-                request_model=request_model,
-                pricing_info=pricing_info,
-                server_port=server_port,
-                server_address=server_address,
-                environment=environment,
-                application_name=application_name,
-                metrics=metrics,
-                event_provider=event_provider,
-                start_time=start_time,
-                span=span,
-                capture_message_content=capture_message_content,
-                disable_metrics=disable_metrics,
-                version=version,
-                **kwargs
-            )
+            except Exception as e:
+                handle_exception(span, e)
 
         return response