Skip to content

Commit b3dbcfa

Browse files
alizenhomalizenhom
committed
litellm kickoff
1 parent 43b12ea commit b3dbcfa

File tree

8 files changed

+308
-5
lines changed

8 files changed

+308
-5
lines changed

‎src/examples/litellm_example/basic.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
from langtrace_python_sdk import with_langtrace_root_span, langtrace
2+
from dotenv import load_dotenv
3+
from litellm import completion, acompletion
4+
import litellm
5+
6+
litellm.set_verbose = False
7+
load_dotenv()
8+
langtrace.init(write_spans_to_console=True)
9+
10+
11+
@with_langtrace_root_span("Litellm Example OpenAI")
12+
def openAI():
13+
response = completion(
14+
model="gpt-3.5-turbo",
15+
messages=[{"content": "Hello, how are you?", "role": "user"}],
16+
)
17+
return response
18+
19+
20+
@with_langtrace_root_span("Litellm Example Anthropic Completion")
21+
def anthropic():
22+
response = completion(
23+
model="claude-2",
24+
messages=[{"content": "Hello, how are you?", "role": "user"}],
25+
temperature=0.5,
26+
top_p=0.5,
27+
n=1,
28+
)
29+
print(response)
30+
return response
31+
32+
33+
@with_langtrace_root_span("Litellm Example Anthropic Streaming")
34+
def anthropic_streaming():
35+
response = completion(
36+
model="claude-2",
37+
messages=[{"content": "Hello, how are you?", "role": "user"}],
38+
stream=True,
39+
temperature=0.5,
40+
# presence_penalty=0.5,
41+
# frequency_penalty=0.5,
42+
top_p=0.5,
43+
n=1,
44+
# logit_bias={"Hello": 1.0},
45+
# top_logprobs=1,
46+
)
47+
for _ in response:
48+
pass
49+
50+
return response
51+
52+
53+
@with_langtrace_root_span("Litellm Example OpenAI Async Streaming")
54+
async def async_anthropic_streaming():
55+
response = await acompletion(
56+
model="claude-2",
57+
messages=[{"content": "Hello, how are you?", "role": "user"}],
58+
stream=True,
59+
temperature=0.5,
60+
top_p=0.5,
61+
n=1,
62+
)
63+
async for _ in response:
64+
pass

‎src/langtrace_python_sdk/constants/instrumentation/common.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
"QDRANT": "Qdrant",
2727
"WEAVIATE": "Weaviate",
2828
"OLLAMA": "Ollama",
29+
"LITELLM": "Litellm",
2930
}
3031

3132
LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

‎src/langtrace_python_sdk/instrumentation/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from .weaviate import WeaviateInstrumentation
1515
from .ollama import OllamaInstrumentor
1616
from .dspy import DspyInstrumentation
17+
from .litellm import LiteLLMInstrumentation
1718

1819
__all__ = [
1920
"AnthropicInstrumentation",
@@ -32,4 +33,5 @@
3233
"WeaviateInstrumentation",
3334
"OllamaInstrumentor",
3435
"DspyInstrumentation",
36+
"LiteLLMInstrumentation",
3537
]

‎src/langtrace_python_sdk/instrumentation/litellm/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
from .instrumentation import LiteLLMInstrumentation
2+
3+
4+
__all__ = ["LiteLLMInstrumentation"]

‎src/langtrace_python_sdk/instrumentation/litellm/instrumentation.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
"""
2+
Copyright (c) 2024 Scale3 Labs
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
"""
16+
17+
from importlib_metadata import version as v
18+
import logging
19+
from typing import Collection
20+
21+
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
22+
from opentelemetry.trace import get_tracer
23+
from wrapt import wrap_function_wrapper as _W
24+
from .patch import litellm_patch, async_litellm_patch
25+
26+
logging.basicConfig(level=logging.FATAL)
27+
28+
29+
class LiteLLMInstrumentation(BaseInstrumentor):
30+
def instrumentation_dependencies(self) -> Collection[str]:
31+
return ["litellm >= 1.0.0"]
32+
33+
def _instrument(self, **kwargs):
34+
modules_to_instrument = [
35+
("litellm.llms.anthropic_text", "AnthropicTextCompletion.completion"),
36+
("litellm.llms.anthropic", "AnthropicChatCompletion.completion"),
37+
]
38+
39+
async_modules_to_instrument = [
40+
("litellm.llms.anthropic_text", "AnthropicTextCompletion.async_completion"),
41+
("litellm.llms.anthropic_text", "AnthropicTextCompletion.async_streaming"),
42+
]
43+
tracer_provider = kwargs.get("tracer_provider")
44+
tracer = get_tracer(__name__, "", tracer_provider)
45+
46+
version = v("litellm")
47+
48+
for module, name in async_modules_to_instrument:
49+
_W(
50+
module=module,
51+
name=name,
52+
wrapper=async_litellm_patch(name, tracer, version),
53+
)
54+
55+
for module, name in modules_to_instrument:
56+
_W(
57+
module=module,
58+
name=name,
59+
wrapper=litellm_patch(name, tracer, version),
60+
)
61+
62+
def _instrument_module(self, module_name):
63+
pass
64+
65+
def _uninstrument(self, **kwargs):
66+
pass

‎src/langtrace_python_sdk/instrumentation/litellm/patch.py‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
2+
from langtrace_python_sdk.instrumentation.openai.patch import StreamWrapper
3+
from langtrace_python_sdk.utils import set_span_attribute
4+
from langtrace_python_sdk.utils.llm import (
5+
calculate_prompt_tokens,
6+
get_extra_attributes,
7+
get_llm_request_attributes,
8+
get_langtrace_attributes,
9+
get_streaming_tokens,
10+
is_streaming,
11+
set_usage_attributes,
12+
)
13+
from langtrace.trace_attributes import LLMSpanAttributes, SpanAttributes, Event
14+
from langtrace_python_sdk.utils.silently_fail import silently_fail
15+
from openai import NOT_GIVEN
16+
from opentelemetry.trace import SpanKind, StatusCode, Status
17+
from opentelemetry import trace
18+
from opentelemetry.trace.propagation import set_span_in_context
19+
20+
import json
21+
22+
23+
def litellm_patch(name, tracer, version):
24+
def traced_method(wrapped, instance, args, kwargs):
25+
26+
print("Name", name)
27+
service_provider = SERVICE_PROVIDERS["LITELLM"]
28+
prompts = kwargs.get("messages")
29+
optional_params = kwargs.get("optional_params", {})
30+
options = {**kwargs, **optional_params}
31+
32+
span_attributes = {
33+
**get_langtrace_attributes(version, service_provider),
34+
**get_llm_request_attributes(options, prompts=prompts),
35+
SpanAttributes.LLM_URL: kwargs.get("api_base"),
36+
SpanAttributes.LLM_PATH: "completion",
37+
**get_extra_attributes(),
38+
}
39+
40+
attributes = LLMSpanAttributes(**span_attributes)
41+
with tracer.start_span(
42+
name=name,
43+
kind=SpanKind.CLIENT,
44+
context=set_span_in_context(trace.get_current_span()),
45+
) as span:
46+
47+
try:
48+
set_input_attributes(span, attributes)
49+
result = wrapped(*args, **kwargs)
50+
if is_streaming(kwargs):
51+
return StreamWrapper(
52+
stream=result,
53+
span=trace.get_current_span(),
54+
prompt_tokens=get_streaming_tokens(kwargs),
55+
function_call=kwargs.get("functions") is not None,
56+
tool_calls=kwargs.get("tools") is not None,
57+
)
58+
59+
else:
60+
print("not streaming", span)
61+
set_response_attributes(span, result)
62+
return result
63+
64+
except Exception as err:
65+
span.record_exception(err)
66+
span.set_status(Status(StatusCode.ERROR, str(err)))
67+
span.end()
68+
raise
69+
70+
return traced_method
71+
72+
73+
def async_litellm_patch(name, tracer, version):
74+
async def traced_method(wrapped, instance, args, kwargs):
75+
print("Name", name)
76+
service_provider = SERVICE_PROVIDERS["LITELLM"]
77+
prompts = kwargs.get("messages")
78+
optional_params = kwargs.get("optional_params", {})
79+
options = {**kwargs, **optional_params}
80+
81+
span_attributes = {
82+
**get_langtrace_attributes(version, service_provider),
83+
**get_llm_request_attributes(options, prompts=prompts),
84+
SpanAttributes.LLM_URL: kwargs.get("api_base"),
85+
SpanAttributes.LLM_PATH: "completion",
86+
**get_extra_attributes(),
87+
}
88+
89+
attributes = LLMSpanAttributes(**span_attributes)
90+
span = tracer.start_span(
91+
name=name,
92+
kind=SpanKind.CLIENT,
93+
context=set_span_in_context(trace.get_current_span()),
94+
)
95+
with tracer.start_span(
96+
name=name,
97+
kind=SpanKind.CLIENT,
98+
context=set_span_in_context(trace.get_current_span()),
99+
) as span:
100+
101+
try:
102+
set_input_attributes(span, attributes)
103+
result = await wrapped(*args, **kwargs)
104+
if is_streaming(kwargs):
105+
return StreamWrapper(
106+
stream=result,
107+
span=span,
108+
prompt_tokens=get_streaming_tokens(kwargs),
109+
function_call=kwargs.get("functions") is not None,
110+
tool_calls=kwargs.get("tools") is not None,
111+
)
112+
113+
else:
114+
set_response_attributes(span, result)
115+
return result
116+
117+
except Exception as err:
118+
span.record_exception(err)
119+
span.set_status(Status(StatusCode.ERROR, str(err)))
120+
span.end()
121+
raise
122+
123+
return traced_method
124+
125+
126+
@silently_fail
127+
def set_input_attributes(span, attributes):
128+
for field, value in attributes.model_dump(by_alias=True).items():
129+
set_span_attribute(span, field, value)
130+
131+
132+
@silently_fail
133+
def set_response_attributes(span, result):
134+
set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, result.model)
135+
set_span_attribute(span, SpanAttributes.LLM_RESPONSE_ID, result.id)
136+
set_span_attribute(
137+
span, SpanAttributes.LLM_SYSTEM_FINGERPRINT, result.system_fingerprint
138+
)
139+
set_usage_attributes(span, result.usage)
140+
141+
for choice in result.choices:
142+
print("Choice", span.is_recording())
143+
set_span_attribute(
144+
span, SpanAttributes.LLM_RESPONSE_FINISH_REASON, choice.finish_reason
145+
)
146+
set_span_attribute(
147+
span, SpanAttributes.LLM_COMPLETIONS, json.dumps(choice.message)
148+
)

‎src/langtrace_python_sdk/langtrace.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
WeaviateInstrumentation,
5454
OllamaInstrumentor,
5555
DspyInstrumentation,
56+
LiteLLMInstrumentation,
5657
)
5758
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
5859
from colorama import Fore
@@ -112,6 +113,7 @@ def init(
112113
"ollama": OllamaInstrumentor(),
113114
"dspy": DspyInstrumentation(),
114115
"crewai": CrewAIInstrumentation(),
116+
"litellm": LiteLLMInstrumentation(),
115117
}
116118

117119
init_instrumentations(disable_instrumentations, all_instrumentations)

‎src/langtrace_python_sdk/utils/llm.py‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,21 @@ def calculate_prompt_tokens(prompt_content, model):
5757
return estimate_tokens(prompt_content) # Fallback method
5858

5959

60+
def get_streaming_tokens(kwargs):
61+
prompt_tokens = 0
62+
for message in kwargs.get("messages", {}):
63+
prompt_tokens += calculate_prompt_tokens(
64+
json.dumps(message), kwargs.get("model")
65+
)
66+
67+
if kwargs.get("functions") is not None and kwargs.get("functions") != NOT_GIVEN:
68+
for function in kwargs.get("functions"):
69+
prompt_tokens += calculate_prompt_tokens(
70+
json.dumps(function), kwargs.get("model")
71+
)
72+
return prompt_tokens
73+
74+
6075
def calculate_price_from_usage(model, usage):
6176
"""
6277
Calculate the price of a model based on its usage."""
@@ -137,11 +152,12 @@ def get_base_url(instance):
137152

138153

139154
def is_streaming(kwargs):
140-
return not (
141-
kwargs.get("stream") is False
142-
or kwargs.get("stream") is None
143-
or kwargs.get("stream") == NOT_GIVEN
144-
)
155+
streaming = kwargs.get("stream", None)
156+
157+
if "optional_params" in kwargs:
158+
streaming = kwargs.get("optional_params").get("stream", None)
159+
160+
return not (streaming is False or streaming is None or streaming == NOT_GIVEN)
145161

146162

147163
def set_usage_attributes(span, usage):

0 commit comments

Comments
 (0)