Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix litellm integration #1726

Merged
merged 2 commits into from
May 31, 2024
Merged

Fix litellm integration #1726

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
02ed5da
- Removed requirement llm
mmabrouk May 31, 2024
404efa9
fix import
mmabrouk May 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions agenta-cli/agenta/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,13 @@

from .sdk.tracing.llm_tracing import Tracing
from .sdk.decorators.tracing import instrument
from .sdk.tracing.callbacks import AgentaLiteLLMHandler
mmabrouk marked this conversation as resolved.
Show resolved Hide resolved
from .sdk.decorators.llm_entrypoint import entrypoint, app
from .sdk.agenta_init import Config, AgentaSingleton, init
from .sdk.utils.helper.openai_cost import calculate_token_usage
from .sdk.client import Agenta
from .sdk.tracing import callbacks


config = PreInitObject("agenta.config", Config)
DEFAULT_AGENTA_SINGLETON_INSTANCE = AgentaSingleton()
tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing # type: ignore
agenta_instrument_handler = AgentaLiteLLMHandler()
3 changes: 0 additions & 3 deletions agenta-cli/agenta/sdk/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,10 @@

from .tracing.llm_tracing import Tracing
from .decorators.tracing import instrument
from .tracing.callbacks import AgentaLiteLLMHandler
from .decorators.llm_entrypoint import entrypoint, app
from .agenta_init import Config, AgentaSingleton, init
from .utils.helper.openai_cost import calculate_token_usage


config = PreInitObject("agenta.config", Config)
DEFAULT_AGENTA_SINGLETON_INSTANCE = AgentaSingleton()
tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing # type: ignore
agenta_instrument_handler = AgentaLiteLLMHandler()
Empty file added agenta-cli/agenta/sdk/tracing/__init__.py
View file
Open in desktop
Empty file.
265 changes: 148 additions & 117 deletions agenta-cli/agenta/sdk/tracing/callbacks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,132 +1,163 @@
# Own Imports
import agenta as ag

# Third Party Imports
from litellm.utils import ModelResponse
from litellm.integrations.custom_logger import CustomLogger as LitellmCustomLogger

def litellm_handler():
try:
from litellm.utils import ModelResponse
from litellm.integrations.custom_logger import (
CustomLogger as LitellmCustomLogger,
)
except ImportError as exc:
raise ImportError(
"The litellm SDK is not installed. Please install it using `pip install litellm`."
) from exc
except Exception as exc:
raise Exception(
"Unexpected error occurred when importing litellm: {}".format(exc)
) from exc

class AgentaLiteLLMHandler(LitellmCustomLogger):
"""This handler is responsible for instrumenting certain events when using litellm to call LLMs.
class LitellmHandler(LitellmCustomLogger):
"""This handler is responsible for instrumenting certain events when using litellm to call LLMs.

Args:
LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
"""
Args:
LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
"""

@property
def _trace(self):
return ag.tracing
@property
def _trace(self):
return ag.tracing

def log_pre_api_call(self, model, messages, kwargs):
self._trace.start_span(
name=kwargs["litellm_params"][
"litellm_call_id"
], # make use of the litellm_call_id as the span name
input=kwargs["input"],
spankind=(
"llm"
if kwargs.get("call_type") in ["completion", "acompletion"]
else (
"embedding"
if kwargs.get("call_type") in ["embedding", "aembedding"]
else "unset"
)
),
)
self._trace.set_span_attribute(
{
"model_config",
def log_pre_api_call(self, model, messages, kwargs):
self._trace.start_span(
name=kwargs["litellm_params"][
"litellm_call_id"
], # make use of the litellm_call_id as the span name
input=kwargs["input"],
spankind=(
"llm"
if kwargs.get("call_type") in ["completion", "acompletion"]
else (
"embedding"
if kwargs.get("call_type") in ["embedding", "aembedding"]
else "unset"
)
),
)
self._trace.set_span_attribute(
{
"model": kwargs.get("model"),
**kwargs.get("optional_params"), # model-specific params passed in
"model_config": {
"model": kwargs.get("model"),
**kwargs.get(
"optional_params"
), # model-specific params passed in
},
}
)

def log_stream_event(self, kwargs, response_obj, start_time, end_time):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": kwargs.get(
"complete_streaming_response"
), # the complete streamed response (only set if `completion(..stream=True)`)
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get(
"response_cost"
), # litellm calculates response cost
},
}
)
)

def log_stream_event(self, kwargs, response_obj, start_time, end_time):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": kwargs.get(
"complete_streaming_response"
), # the complete streamed response (only set if `completion(..stream=True)`)
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get("response_cost"), # litellm calculates response cost
},
)
def log_success_event(
self, kwargs, response_obj: ModelResponse, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": response_obj.choices[0].message.content,
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get(
"response_cost"
), # litellm calculates response cost
},
)

def log_success_event(
self, kwargs, response_obj: ModelResponse, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": response_obj.choices[0].message.content,
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get("response_cost"), # litellm calculates response cost
},
)
def log_failure_event(
self, kwargs, response_obj: ModelResponse, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
self._trace.set_span_attribute(
attributes={
"traceback_exception": kwargs[
"traceback_exception"
], # the traceback generated via `traceback.format_exc()`
"call_end_time": kwargs[
"end_time"
], # datetime object of when call was completed
},
)
self._trace.end_span(
outputs={
"message": kwargs["exception"], # the Exception raised
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get(
"response_cost"
), # litellm calculates response cost
},
)

def log_failure_event(
self, kwargs, response_obj: ModelResponse, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
self._trace.set_span_attribute(
attributes={
"traceback_exception": kwargs[
"traceback_exception"
], # the traceback generated via `traceback.format_exc()`
"call_end_time": kwargs[
"end_time"
], # datetime object of when call was completed
},
)
self._trace.end_span(
outputs={
"message": kwargs["exception"], # the Exception raised
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get("response_cost"), # litellm calculates response cost
},
)
async def async_log_stream_event(
self, kwargs, response_obj, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": kwargs.get(
"complete_streaming_response"
), # the complete streamed response (only set if `completion(..stream=True)`)
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get(
"response_cost"
), # litellm calculates response cost
},
)

async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": kwargs.get(
"complete_streaming_response"
), # the complete streamed response (only set if `completion(..stream=True)`)
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get("response_cost"), # litellm calculates response cost
},
)
async def async_log_success_event(
self, kwargs, response_obj, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": response_obj.choices[0].message.content,
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get(
"response_cost"
), # litellm calculates response cost
},
)

async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self._trace.update_span_status(span=self._trace.active_span, value="OK")
self._trace.end_span(
outputs={
"message": response_obj.choices[0].message.content,
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get("response_cost"), # litellm calculates response cost
},
)
async def async_log_failure_event(
self, kwargs, response_obj, start_time, end_time
):
self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
self._trace.set_span_attribute(
attributes={
"traceback_exception": kwargs[
"traceback_exception"
], # the traceback generated via `traceback.format_exc()`
"call_end_time": kwargs[
"end_time"
], # datetime object of when call was completed
},
)
self._trace.end_span(
outputs={
"message": kwargs["exception"], # the Exception raised
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get(
"response_cost"
), # litellm calculates response cost
},
)

async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
self._trace.set_span_attribute(
attributes={
"traceback_exception": kwargs[
"traceback_exception"
], # the traceback generated via `traceback.format_exc()`
"call_end_time": kwargs[
"end_time"
], # datetime object of when call was completed
},
)
self._trace.end_span(
outputs={
"message": kwargs["exception"], # the Exception raised
"usage": kwargs.get("usage"), # litellm calculates usage
"cost": kwargs.get("response_cost"), # litellm calculates response cost
},
)
return LitellmHandler()
Loading
Loading