|
4 | 4 | import os
|
5 | 5 | import json
|
6 | 6 | import time
|
| 7 | +import sys |
7 | 8 | from typing import Collection
|
8 | 9 | from opentelemetry.instrumentation.ollama.config import Config
|
9 | 10 | from opentelemetry.instrumentation.ollama.utils import dont_throw
|
@@ -431,6 +432,18 @@ def is_metrics_collection_enabled() -> bool:
|
431 | 432 | return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true"
|
432 | 433 |
|
433 | 434 |
|
| 435 | +# Proxy class to dynamically call the latest ollama function implementations |
| 436 | +class InstrumentedFunction: |
| 437 | + """Proxy function that always invokes the latest ollama function implementation""" |
| 438 | + def __init__(self, func_name): |
| 439 | + self.func_name = func_name |
| 440 | + |
| 441 | + def __call__(self, *args, **kwargs): |
| 442 | + import ollama |
| 443 | + actual_func = getattr(ollama, self.func_name) |
| 444 | + return actual_func(*args, **kwargs) |
| 445 | + |
| 446 | + |
434 | 447 | class OllamaInstrumentor(BaseInstrumentor):
|
435 | 448 | """An instrumentor for Ollama's client library."""
|
436 | 449 |
|
@@ -477,6 +490,24 @@ def _instrument(self, **kwargs):
|
477 | 490 | _wrap(tracer, token_histogram, duration_histogram, wrapped_method),
|
478 | 491 | )
|
479 | 492 |
|
| 493 | + # replace imported ollama functions in other modules |
| 494 | + try: |
| 495 | + # Iterate through all loaded modules to find modules that reference imported ollama functions |
| 496 | + for module_name, module in list(sys.modules.items()): |
| 497 | + if module_name != "ollama" and hasattr(module, "__dict__"): |
| 498 | + for attr_name, attr_value in list(module.__dict__.items()): |
| 499 | + for wrapped_method in WRAPPED_METHODS: |
| 500 | + method_name = wrapped_method.get("method") |
| 501 | + if ( |
| 502 | + attr_name == method_name |
| 503 | + and hasattr(attr_value, "__module__") |
| 504 | + and attr_value.__module__.startswith("ollama") |
| 505 | + ): |
| 506 | + module.__dict__[attr_name] = InstrumentedFunction(method_name) |
| 507 | + logger.debug(f"Replaced {method_name} in module {module_name}") |
| 508 | + except Exception as e: |
| 509 | + logger.warning(f"Error instrumenting imported ollama methods: {e}") |
| 510 | + |
480 | 511 | def _uninstrument(self, **kwargs):
|
481 | 512 | for wrapped_method in WRAPPED_METHODS:
|
482 | 513 | unwrap(
|
|
0 commit comments