Skip to content

Commit

Permalink
fix(langchain): support custom models (#706)
Browse files Browse the repository at this point in the history
  • Loading branch information
nirga authored Mar 27, 2024
1 parent 4011d3f commit df89031
Show file tree
Hide file tree
Showing 7 changed files with 331 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
{
Expand Down Expand Up @@ -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,
},
]


Expand Down
Original file line number Diff line number Diff line change
@@ -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,
)
Original file line number Diff line number Diff line change
@@ -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."""

Expand All @@ -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")
103 changes: 102 additions & 1 deletion packages/opentelemetry-instrumentation-langchain/poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit df89031

Please sign in to comment.