Add patch_abc utility

Author: pmcollins

docs/examples/patch_abc/langchain_example.py

@@ -0,0 +1,93 @@
+import argparse
+
+from langchain_core.language_models import BaseLLM, BaseChatModel
+from langchain_core.language_models.base import BaseLanguageModel
+
+# important note: if you import these after patching, the patch won't apply!
+# mitigation will be added to patch_abc in a future release
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_ollama import OllamaLLM
+from langchain_openai import ChatOpenAI
+
+from opentelemetry.util._wrap import patch_abc
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="LangChain model comparison")
+    parser.add_argument("--provider", choices=["ollama", "openai", "huggingface"], default="ollama",
+                        help="Choose model provider (default: ollama)")
+    parser.add_argument("--model", type=str, help="Specify model name")
+    parser.add_argument("--prompt", type=str, default="What is the capital of France?",
+                        help="Input prompt")
+
+    return parser.parse_args()
+
+
+def chat_with_model(model: BaseLanguageModel, prompt: str) -> str:
+    try:
+        response = model.invoke(prompt)
+        if hasattr(response, 'content'):
+            return response.content
+        else:
+            return str(response)
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+def create_huggingface_model(model: str = "google/flan-t5-small"):
+    return HuggingFaceEndpoint(
+        repo_id=model,
+        temperature=0.7
+    )
+
+
+def create_openai_model(model: str = "gpt-3.5-turbo"):
+    return ChatOpenAI(
+        model=model,
+        temperature=0.7
+    )
+
+
+def create_ollama_model(model: str = "llama2"):
+    return OllamaLLM(
+        model=model,
+        temperature=0.7
+    )
+
+
+def patch_llm():
+    def my_wrapper(orig_fcn):
+        def wrapped_fcn(self, *args, **kwargs):
+            print("wrapper starting")
+            print(f"Arguments: {args}")
+            print(f"Keyword arguments: {kwargs}")
+            return orig_fcn(self, *args, **kwargs)
+
+        return wrapped_fcn
+
+    patch_abc(BaseLLM, "_generate", my_wrapper)
+
+    # this is for OpenAI, which is a weird case. The _generate method is in a differnt base class and gets called twice.
+    patch_abc(BaseChatModel, "_generate", my_wrapper)
+
+
+def main():
+    args = parse_args()
+
+    patch_llm()
+
+    if args.provider == "ollama":
+        model = create_ollama_model(args.model or "llama2")
+    elif args.provider == "openai":
+        model = create_openai_model(args.model or "gpt-3.5-turbo")
+    elif args.provider == "huggingface":
+        model = create_huggingface_model(args.model or "google/flan-t5-small")
+    else:
+        raise ValueError(f"Unsupported provider: {args.provider}")
+
+    response = chat_with_model(model, args.prompt)
+    print(f"{args.provider.title()} Response: {response}")
+
+
+if __name__ == "__main__":
+    main()

docs/examples/patch_abc/requirements.txt

@@ -0,0 +1,6 @@
+langchain-core
+langchain-openai
+langchain-ollama
+langchain-huggingface
+huggingface-hub
+opentelemetry-api

docs/examples/patch_abc/simple_example.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+
+from opentelemetry.util._wrap import patch_abc
+
+
+class Greeter(ABC):
+    @abstractmethod
+    def greet(self):
+        pass
+
+
+class EngishGreeter(Greeter):
+    def greet(self):
+        print("hello")
+
+
+class SpanishGreeter(Greeter):
+    def greet(self):
+        print("hola")
+
+
+if __name__ == '__main__':
+    def my_wrapper(orig_fcn):
+        def wrapped_fcn(self, *args, **kwargs):
+            print("wrapper running")
+            result = orig_fcn(self, *args, **kwargs)
+            return result
+
+        return wrapped_fcn
+
+
+    patch_abc(Greeter, "greet", my_wrapper)
+
+    EngishGreeter().greet()
+    SpanishGreeter().greet()

opentelemetry-api/src/opentelemetry/util/_wrap.py

@@ -0,0 +1,44 @@
+def patch_abc(abstract_base_class, method_name, w):
+    """
+    Patches a method on leaf subclasses of an abstract base class.
+
+    """
+    all_subclasses = recursively_get_all_subclasses(abstract_base_class)
+    leaf_subclasses = get_leaf_subclasses(all_subclasses)
+
+    for subclass in leaf_subclasses:
+        # Patch if the subclass has the method (either defined or inherited)
+        # and it's actually callable
+        if hasattr(subclass, method_name) and callable(getattr(subclass, method_name)):
+            old_method = getattr(subclass, method_name)
+            setattr(subclass, method_name, w(old_method))
+
+    # This implementation does not work if the instrumented class is imported after the instrumentor runs.
+    # However, that case can be handled by querying the gc module for all existing classes; this capability can be added
+    # in a follow-up change.
+
+
+def get_leaf_subclasses(all_subclasses):
+    """
+    Returns only the leaf classes (classes with no subclasses) from a set of classes.
+    """
+    leaf_classes = set()
+    for cls in all_subclasses:
+        # A class is a leaf if no other class in the set is its subclass
+        is_leaf = True
+        for other_cls in all_subclasses:
+            if other_cls != cls and issubclass(other_cls, cls):
+                is_leaf = False
+                break
+        if is_leaf:
+            leaf_classes.add(cls)
+    return leaf_classes
+
+
+
+def recursively_get_all_subclasses(cls):
+    out = set()
+    for subclass in cls.__subclasses__():
+        out.add(subclass)
+        out.update(recursively_get_all_subclasses(subclass))
+    return out