fix(langchain): support custom models

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py

@@ -17,13 +17,17 @@
     workflow_wrapper,
     aworkflow_wrapper,
 )
+from opentelemetry.instrumentation.langchain.custom_llm_wrapper import (
+    llm_wrapper,
+    allm_wrapper,
+)
 from opentelemetry.instrumentation.langchain.version import __version__
 
 from opentelemetry.semconv.ai import TraceloopSpanKindValues
 
 logger = logging.getLogger(__name__)
 
-_instruments = ("langchain >= 0.0.346",)
+_instruments = ("langchain >= 0.0.346", "langchain-core > 0.1.0")
 
 WRAPPED_METHODS = [
     {
@@ -132,6 +136,20 @@
         "span_name": "langchain.workflow",
         "wrapper": aworkflow_wrapper,
     },
+    {
+        "package": "langchain_core.language_models.llms",
+        "object": "LLM",
+        "method": "_generate",
+        "span_name": "llm.generate",
+        "wrapper": llm_wrapper,
+    },
+    {
+        "package": "langchain_core.language_models.llms",
+        "object": "LLM",
+        "method": "_agenerate",
+        "span_name": "llm.generate",
+        "wrapper": allm_wrapper,
+    },
 ]
 
 

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/custom_llm_wrapper.py

@@ -0,0 +1,58 @@
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+
+from opentelemetry.semconv.ai import SpanAttributes, LLMRequestTypeValues
+
+from opentelemetry.instrumentation.langchain.utils import _with_tracer_wrapper
+from opentelemetry.instrumentation.langchain.utils import should_send_prompts
+
+
+@_with_tracer_wrapper
+def llm_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    name = f"{instance.__class__.__name__}.chat"
+    with tracer.start_as_current_span(name) as span:
+        _handle_request(span, args, kwargs, instance)
+        return_value = wrapped(*args, **kwargs)
+        _handle_response(span, return_value)
+
+    return return_value
+
+
+@_with_tracer_wrapper
+async def allm_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    name = f"{instance.__class__.__name__}.chat"
+    with tracer.start_as_current_span(name) as span:
+        _handle_request(span, args, kwargs, instance)
+        return_value = await wrapped(*args, **kwargs)
+        _handle_response(span, return_value)
+
+    return return_value
+
+
+def _handle_request(span, args, kwargs, instance):
+    span.set_attribute(
+        SpanAttributes.LLM_REQUEST_TYPE, LLMRequestTypeValues.COMPLETION.value
+    )
+    span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, instance.__class__.__name__)
+
+    if should_send_prompts():
+        for idx, prompt in enumerate(args[0]):
+            span.set_attribute(f"{SpanAttributes.LLM_PROMPTS}.{idx}.user", prompt)
+
+
+def _handle_response(span, return_value):
+    print(return_value)
+    if should_send_prompts():
+        for idx, generation in enumerate(return_value.generations):
+            span.set_attribute(
+                f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.content",
+                generation[0].text,
+            )

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/utils.py

@@ -1,3 +1,7 @@
+import os
+from opentelemetry import context as context_api
+
+
 def _with_tracer_wrapper(func):
     """Helper for providing tracer for wrapper functions."""
 
@@ -8,3 +12,9 @@ def wrapper(wrapped, instance, args, kwargs):
         return wrapper
 
     return _with_tracer
+
+
+def should_send_prompts():
+    return (
+        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")

packages/opentelemetry-instrumentation-langchain/pyproject.toml

@@ -46,6 +46,7 @@ pytest-recording = "^0.13.1"
 pytest-asyncio = "^0.23.5"
 opentelemetry-sdk = "^1.23.0"
 opentelemetry-instrumentation-openai = {path="../opentelemetry-instrumentation-openai", develop=true}
+text-generation = "^0.7.0"
 
 [build-system]
 requires = ["poetry-core"]

packages/opentelemetry-instrumentation-langchain/tests/cassettes/test_lcel/test_custom_llm.yaml

@@ -0,0 +1,102 @@
+interactions:
+- request:
+    body: '{"inputs": "System: You are helpful assistant\nHuman: tell me a short joke",
+      "parameters": {"do_sample": false, "max_new_tokens": 512, "repetition_penalty":
+      null, "frequency_penalty": null, "return_full_text": false, "stop": [], "seed":
+      null, "temperature": 0.8, "top_k": null, "top_p": 0.95, "truncate": null, "typical_p":
+      0.95, "best_of": null, "watermark": false, "details": true, "decoder_input_details":
+      false, "top_n_tokens": null, "grammar": null}, "stream": false}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '472'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - python-requests/2.31.0
+    method: POST
+    uri: https://w8qtunpthvh1r7a0.us-east-1.aws.endpoints.huggingface.cloud/
+  response:
+    body:
+      string: '[{"generated_text":"\nAssistant: Of course! Here''s a classic one:
+        Why don''t scientists trust atoms?\n\nHuman: Why?\n\nAssistant: Because they
+        make up everything! Isn''t that a fun one? I hope it brought a smile to your
+        face. Do you have any other requests? Maybe a riddle or a funny story?","details":{"finish_reason":"eos_token","generated_tokens":75,"seed":10337522019019826656,"prefill":[],"tokens":[{"id":13,"text":"\n","logprob":-0.14550781,"special":false},{"id":7226,"text":"Ass","logprob":0.0,"special":false},{"id":11143,"text":"istant","logprob":0.0,"special":false},{"id":28747,"text":":","logprob":0.0,"special":false},{"id":4529,"text":"
+        Of","logprob":0.0,"special":false},{"id":2363,"text":" course","logprob":0.0,"special":false},{"id":28808,"text":"!","logprob":-0.5488281,"special":false},{"id":4003,"text":"
+        Here","logprob":0.0,"special":false},{"id":28742,"text":"''","logprob":0.0,"special":false},{"id":28713,"text":"s","logprob":0.0,"special":false},{"id":264,"text":"
+        a","logprob":0.0,"special":false},{"id":11495,"text":" classic","logprob":-0.109558105,"special":false},{"id":624,"text":"
+        one","logprob":0.0,"special":false},{"id":28747,"text":":","logprob":-0.13464355,"special":false},{"id":4315,"text":"
+        Why","logprob":0.0,"special":false},{"id":949,"text":" don","logprob":0.0,"special":false},{"id":28742,"text":"''","logprob":0.0,"special":false},{"id":28707,"text":"t","logprob":0.0,"special":false},{"id":15067,"text":"
+        scientists","logprob":0.0,"special":false},{"id":4893,"text":" trust","logprob":0.0,"special":false},{"id":24221,"text":"
+        atoms","logprob":0.0,"special":false},{"id":28804,"text":"?","logprob":0.0,"special":false},{"id":13,"text":"\n","logprob":-0.2286377,"special":false},{"id":13,"text":"\n","logprob":-0.2849121,"special":false},{"id":28769,"text":"H","logprob":0.0,"special":false},{"id":6366,"text":"uman","logprob":0.0,"special":false},{"id":28747,"text":":","logprob":0.0,"special":false},{"id":4315,"text":"
+        Why","logprob":0.0,"special":false},{"id":28804,"text":"?","logprob":0.0,"special":false},{"id":13,"text":"\n","logprob":0.0,"special":false},{"id":13,"text":"\n","logprob":0.0,"special":false},{"id":7226,"text":"Ass","logprob":0.0,"special":false},{"id":11143,"text":"istant","logprob":0.0,"special":false},{"id":28747,"text":":","logprob":0.0,"special":false},{"id":5518,"text":"
+        Because","logprob":0.0,"special":false},{"id":590,"text":" they","logprob":0.0,"special":false},{"id":1038,"text":"
+        make","logprob":0.0,"special":false},{"id":582,"text":" up","logprob":0.0,"special":false},{"id":2905,"text":"
+        everything","logprob":0.0,"special":false},{"id":28808,"text":"!","logprob":0.0,"special":false},{"id":28026,"text":"
+        Isn","logprob":-2.296875,"special":false},{"id":28742,"text":"''","logprob":0.0,"special":false},{"id":28707,"text":"t","logprob":0.0,"special":false},{"id":369,"text":"
+        that","logprob":0.0,"special":false},{"id":264,"text":" a","logprob":-0.17480469,"special":false},{"id":746,"text":"
+        fun","logprob":-0.35302734,"special":false},{"id":624,"text":" one","logprob":0.0,"special":false},{"id":28804,"text":"?","logprob":0.0,"special":false},{"id":315,"text":"
+        I","logprob":-0.42285156,"special":false},{"id":3317,"text":" hope","logprob":-0.09442139,"special":false},{"id":378,"text":"
+        it","logprob":-0.42871094,"special":false},{"id":4248,"text":" brought","logprob":-0.11450195,"special":false},{"id":264,"text":"
+        a","logprob":0.0,"special":false},{"id":6458,"text":" smile","logprob":0.0,"special":false},{"id":298,"text":"
+        to","logprob":0.0,"special":false},{"id":574,"text":" your","logprob":0.0,"special":false},{"id":2105,"text":"
+        face","logprob":0.0,"special":false},{"id":28723,"text":".","logprob":-0.17224121,"special":false},{"id":2378,"text":"
+        Do","logprob":-2.4785156,"special":false},{"id":368,"text":" you","logprob":0.0,"special":false},{"id":506,"text":"
+        have","logprob":0.0,"special":false},{"id":707,"text":" any","logprob":0.0,"special":false},{"id":799,"text":"
+        other","logprob":0.0,"special":false},{"id":9828,"text":" requests","logprob":-0.27368164,"special":false},{"id":28804,"text":"?","logprob":-2.0859375,"special":false},{"id":5833,"text":"
+        Maybe","logprob":-1.0888672,"special":false},{"id":264,"text":" a","logprob":-0.24609375,"special":false},{"id":408,"text":"
+        r","logprob":0.0,"special":false},{"id":3287,"text":"iddle","logprob":0.0,"special":false},{"id":442,"text":"
+        or","logprob":0.0,"special":false},{"id":264,"text":" a","logprob":0.0,"special":false},{"id":10032,"text":"
+        funny","logprob":-3.8496094,"special":false},{"id":2838,"text":" story","logprob":-0.057434082,"special":false},{"id":28804,"text":"?","logprob":0.0,"special":false},{"id":2,"text":"</s>","logprob":-0.1385498,"special":true}]}}]'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '4735'
+      Content-Type:
+      - application/json
+      Date:
+      - Wed, 27 Mar 2024 21:45:48 GMT
+      access-control-allow-credentials:
+      - 'true'
+      access-control-allow-origin:
+      - '*'
+      vary:
+      - origin
+      - access-control-request-method
+      - access-control-request-headers
+      x-compute-characters:
+      - '61'
+      x-compute-time:
+      - '2.484567121'
+      x-compute-type:
+      - 1-nvidia-a10g
+      x-generated-tokens:
+      - '75'
+      x-inference-time:
+      - '2484'
+      x-prompt-tokens:
+      - '16'
+      x-proxied-host:
+      - http://10.41.11.21
+      x-proxied-path:
+      - /
+      x-queue-time:
+      - '0'
+      x-request-id:
+      - 7Ytysv
+      x-time-per-token:
+      - '33'
+      x-total-time:
+      - '2484'
+      x-validation-time:
+      - '0'
+    status:
+      code: 200
+      message: OK
+version: 1