diff --git a/agentops/client.py b/agentops/client.py index c082642c..5447f510 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -255,7 +255,7 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi self._session = None return logger.warning("šŸ–‡ AgentOps: Cannot start session") - logger.info('View info on this session at https://app.agentops.ai/drilldown?session_id=%s', + logger.info('View info on this session at https://app.agentops.ai/drilldown?session_id=%s', self._session.session_id) return self._session.session_id @@ -277,19 +277,19 @@ def end_session(self, if not any(end_state == state.value for state in EndState): return logger.warning("šŸ–‡ AgentOps: Invalid end_state. Please use one of the EndState enums") - + if self._worker is None or self._worker._session is None: return logger.warning("šŸ–‡ AgentOps: Cannot end session - no current worker or session") self._session.video = video self._session.end_session(end_state, end_state_reason) token_cost = self._worker.end_session(self._session) - + if token_cost == 'unknown': print('šŸ–‡ AgentOps: Could not determine cost of run.') else: token_cost_d = Decimal(token_cost) - print('šŸ–‡ AgentOps: This run cost ${}'.format('{:.2f}'.format( + print('\nšŸ–‡ AgentOps: This run cost ${}'.format('{:.2f}'.format( token_cost_d) if token_cost_d == 0 else '{:.6f}'.format(token_cost_d))) self._session = None self._worker = None diff --git a/agentops/llm_tracker.py b/agentops/llm_tracker.py index 021d6fef..b2e411db 100644 --- a/agentops/llm_tracker.py +++ b/agentops/llm_tracker.py @@ -4,7 +4,7 @@ from importlib.metadata import version from packaging.version import Version, parse from .log_config import logger -from .event import LLMEvent, ToolEvent, ErrorEvent +from .event import LLMEvent, ErrorEvent from .helpers import get_ISO_time, check_call_stack_for_agent_id import inspect from typing import Optional @@ -28,8 +28,9 @@ class LlmTracker: ), }, 'cohere': { - '1.0.0': ( + '5.4.0': ( "chat", + "chat_stream" ), } } @@ -474,11 +475,14 @@ def override_api(self): module = import_module(api) if api == 'litellm': module_version = version(api) + if module_version is None: + logger.warning(f'šŸ–‡ AgentOps: Cannot determine LiteLLM version. Only LiteLLM>=1.3.1 supported.') + if Version(module_version) >= parse('1.3.1'): self.override_litellm_completion() self.override_litellm_async_completion() else: - logger.warning(f'šŸ–‡ AgentOps: Only litellm>=1.3.1 supported. v{module_version} found.') + logger.warning(f'šŸ–‡ AgentOps: Only LiteLLM>=1.3.1 supported. v{module_version} found.') return # If using an abstraction like litellm, do not patch the underlying LLM APIs if api == 'openai': @@ -494,14 +498,16 @@ def override_api(self): self._override_method(api, method_path, module) if api == 'cohere': - # Patch openai v1.0.0+ methods - if hasattr(module, '__version__'): - module_version = parse(module.__version__) - if True: # TODO: check version - self.override_cohere_chat() - self.override_cohere_chat_stream() + # Patch cohere vx.x.x+ methods + module_version = version(api) + if module_version is None: + logger.warning(f'šŸ–‡ AgentOps: Cannot determine Cohere version. Only Cohere>=5.4.0 supported.') + + if Version(module_version) >= parse('5.4.0'): + self.override_cohere_chat() + self.override_cohere_chat_stream() else: - logger.warning(f'šŸ–‡ AgentOps: Only Cohere>=x.x.x supported. v{module_version} found.') + logger.warning(f'šŸ–‡ AgentOps: Only Cohere>=5.4.0 supported. v{module_version} found.') def stop_instrumenting(self): self.undo_override_openai_v1_async_completion() diff --git a/tests/test_cohere.py b/tests/test_cohere.py index 27581a9e..a781d944 100644 --- a/tests/test_cohere.py +++ b/tests/test_cohere.py @@ -15,23 +15,18 @@ }, ], message="What year was he born?", - model="command" + model="command", + connectors=[{"id": "web-search"}] ) print(chat) -# import cohere -# import agentops -# from dotenv import load_dotenv -# load_dotenv() - -# agentops.init(endpoint='http://localhost:8000') -# co = cohere.Client() +stream = co.chat_stream( + message="Tell me a short story" +) -# stream = co.chat_stream( -# message="Tell me a short story" -# ) +for event in stream: + if event.event_type == "text-generation": + print(event.text, end='') -# for event in stream: -# if event.event_type == "text-generation": -# print(event.text, end='') +agentops.end_session('Success')