Done

AgentOps-AI · May 10, 2024 · 0a2517f · 0a2517f
1 parent e83c62c
commit 0a2517f
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 28 deletions.
diff --git 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
@@ -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
@@ -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')