Merge branch 'main' into v2-api

agentops/agent.py

@@ -18,17 +18,17 @@ def new_init(self, *args, **kwargs):
                 try:
                     original_init(self, *args, **kwargs)
                     self.agent_ops_agent_id = str(uuid4())
-                    Client().create_agent(self.agent_ops_agent_id, self.agent_ops_agent_name)
+                    Client().create_agent(name=self.agent_ops_agent_name, agent_id=self.agent_ops_agent_id)
                 except AttributeError as e:
                     logger.warning("AgentOps failed to track an agent. This often happens if agentops.init() was not "
-                                  "called before initializing an agent with the @track_agent decorator.")
+                                   "called before initializing an agent with the @track_agent decorator.")
                     raise e
 
             obj.__init__ = new_init
 
         elif isfunction(obj):
             obj.agent_ops_agent_id = str(uuid4())
-            Client().create_agent(obj.agent_ops_agent_id, obj.agent_ops_agent_name)
+            Client().create_agent(name=obj.agent_ops_agent_name, agent_id=obj.agent_ops_agent_id)
 
         else:
             raise Exception("Invalid input, 'obj' must be a class or a function")

agentops/client.py

@@ -83,7 +83,6 @@ def __init__(self,
         self._tags: Optional[List[str]] = tags
         self._tags_for_future_session: Optional[List[str]] = None
 
-
         self._env_data_opt_out = os.getenv('AGENTOPS_ENV_DATA_OPT_OUT') and os.getenv(
             'AGENTOPS_ENV_DATA_OPT_OUT').lower() == 'true'
 
@@ -98,7 +97,8 @@ def __init__(self,
 
         self._handle_unclean_exits()
 
-        instrument_llm_calls, auto_start_session = self._check_for_partner_frameworks(instrument_llm_calls, auto_start_session)
+        instrument_llm_calls, auto_start_session = self._check_for_partner_frameworks(
+            instrument_llm_calls, auto_start_session)
 
         if auto_start_session:
             self.start_session(tags, self.config, inherited_session_id)
@@ -114,11 +114,16 @@ def _check_for_partner_frameworks(self, instrument_llm_calls, auto_start_session
         for framework in partner_frameworks.keys():
             if framework in sys.modules:
                 self.add_tags([framework])
-                if 'autogen':
-                    import autogen
-                    autogen.runtime_logging.start(logger_type="agentops")
+                if framework == 'autogen':
+                    try:
+                        import autogen
+                        autogen.runtime_logging.start(logger_type="agentops")
+                    except ImportError:
+                        pass
+                    except Exception as e:
+                        logger.warning("🖇️ AgentOps: Failed to set up autogen logger with AgentOps. Error: " + e)
 
-                return partner_frameworks[framework]
+                    return partner_frameworks[framework]
 
         return instrument_llm_calls, auto_start_session
 
@@ -334,7 +339,7 @@ def create_agent(self, name: str, agent_id: Optional[str] = None) -> str:
         if agent_id is None:
             agent_id = str(uuid.uuid4())
         if self._worker:
-            self._worker.create_agent(agent_id, name)
+            self._worker.create_agent(name=name, agent_id=agent_id)
             return agent_id
 
     def _handle_unclean_exits(self):
@@ -408,4 +413,5 @@ def parent_key(self):
         return self.config.parent_key
 
     def stop_instrumenting(self):
-        self.llm_tracker.stop_instrumenting()
+        if self.llm_tracker:
+            self.llm_tracker.stop_instrumenting()

agentops/host_env.py

@@ -2,13 +2,78 @@
 import psutil
 import socket
 from .helpers import get_agentops_version
+import importlib.metadata
+import os
+import sys
 
 
 def get_sdk_details():
     try:
         return {
             "AgentOps SDK Version": get_agentops_version(),
             "Python Version": platform.python_version(),
+            "System Packages": get_sys_packages()
+        }
+    except:
+        return {}
+
+
+def get_python_details():
+    try:
+        return {
+            "Python Version": platform.python_version()
+        }
+    except:
+        return {}
+
+
+def get_agentops_details():
+    try:
+        return {
+            "AgentOps SDK Version": get_agentops_version()
+        }
+    except:
+        return {}
+
+
+def get_sys_packages():
+    sys_packages = {}
+    for module in sys.modules:
+        try:
+            version = importlib.metadata.version(module)
+            sys_packages[module] = version
+        except importlib.metadata.PackageNotFoundError:
+            # Skip built-in modules and those without package metadata
+            continue
+
+    return sys_packages
+
+
+def get_installed_packages():
+
+    try:
+        return {
+            # TODO: test
+            # TODO: add to opt out
+            "Installed Packages": {dist.metadata['Name']: dist.version for dist in importlib.metadata.distributions()}
+        }
+    except:
+        return {}
+
+
+def get_current_directory():
+    try:
+        return {
+            "Project Working Directory": os.getcwd()
+        }
+    except:
+        return {}
+
+
+def get_virtual_env():
+    try:
+        return {
+            "Virtual Environment": os.environ.get('VIRTUAL_ENV', None)
         }
     except:
         return {}
@@ -51,7 +116,6 @@ def get_ram_details():
         return {}
 
 
-
 def get_disk_details():
     partitions = psutil.disk_partitions()
     disk_info = {}
@@ -71,7 +135,9 @@ def get_host_env(opt_out: bool = False):
     if opt_out:
         return {
             "SDK": get_sdk_details(),
-            "OS": get_os_details()
+            "OS": get_os_details(),
+            "Project Working Directory": get_current_directory(),
+            "Virtual Environment": get_virtual_env()
         }
     else:
         return {
@@ -80,4 +146,7 @@ def get_host_env(opt_out: bool = False):
             "CPU": get_cpu_details(),
             "RAM": get_ram_details(),
             "Disk": get_disk_details(),
+            "Installed Packages": get_installed_packages(),
+            "Project Working Directory": get_current_directory(),
+            "Virtual Environment": get_virtual_env()
         }

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, ErrorEvent
+from .event import LLMEvent, ActionEvent, ToolEvent, ErrorEvent
 from .helpers import get_ISO_time, check_call_stack_for_agent_id
 import inspect
 from typing import Optional
@@ -228,7 +228,7 @@ async def async_generator():
         return response
 
     def _handle_response_cohere(self, response, kwargs, init_timestamp):
-        # TODO: """Handle responses for Cohere versions >v5.4.0"""
+        """Handle responses for Cohere versions >v5.4.0"""
         from cohere.types.non_streamed_chat_response import NonStreamedChatResponse
         from cohere.types.streamed_chat_response import (
             StreamedChatResponse,
@@ -248,6 +248,8 @@ def _handle_response_cohere(self, response, kwargs, init_timestamp):
             params=kwargs
         )
 
+        self.action_events = {}
+
         def handle_stream_chunk(chunk):
 
             # We take the first chunk and accumulate the deltas from all subsequent chunks to build one full chat completion
@@ -261,21 +263,59 @@ def handle_stream_chunk(chunk):
 
             try:
                 if isinstance(chunk, StreamedChatResponse_StreamEnd):
-                    # Streaming is done. Record LLMEvent
-                    # self.llm_event.returns.finish_reason = chunk.is_finished
+                    # StreamedChatResponse_TextGeneration = LLMEvent
                     self.llm_event.completion = {
-                        "role": "assistant", "content": self.llm_event.completion}
+                        "role": "assistant", "content": chunk.response.text}
                     self.llm_event.end_timestamp = get_ISO_time()
-
                     self.client.record(self.llm_event)
 
+                    # StreamedChatResponse_SearchResults = ActionEvent
+                    search_results = chunk.response.search_results
+                    for search_result in search_results:
+                        query = search_result.search_query
+                        if query.generation_id in self.action_events:
+                            action_event = self.action_events[query.generation_id]
+                            search_result_dict = search_result.dict()
+                            del search_result_dict["search_query"]
+                            action_event.returns = search_result_dict
+                            action_event.end_timestamp = get_ISO_time()
+
+                    # StreamedChatResponse_CitationGeneration = ActionEvent
+                    documents = {doc['id']: doc for doc in chunk.response.documents}
+                    citations = chunk.response.citations
+                    for citation in citations:
+                        citation_id = f"{citation.start}.{citation.end}"
+                        if citation_id in self.action_events:
+                            action_event = self.action_events[citation_id]
+                            citation_dict = citation.dict()
+                            # Replace document_ids with the actual documents
+                            citation_dict['documents'] = [documents[doc_id]
+                                                          for doc_id in citation_dict['document_ids'] if doc_id in documents]
+                            del citation_dict['document_ids']
+
+                            action_event.returns = citation_dict
+                            action_event.end_timestamp = get_ISO_time()
+
+                    for key, action_event in self.action_events.items():
+                        self.client.record(action_event)
+
                 elif isinstance(chunk, StreamedChatResponse_TextGeneration):
                     self.llm_event.completion += chunk.text
                 elif isinstance(chunk, StreamedChatResponse_ToolCallsGeneration):
                     pass
                 elif isinstance(chunk, StreamedChatResponse_CitationGeneration):
-                    pass
+                    for citation in chunk.citations:
+                        self.action_events[f"{citation.start}.{citation.end}"] = ActionEvent(
+                            action_type="citation",
+                            init_timestamp=get_ISO_time(),
+                            params=citation.text)
                 elif isinstance(chunk, StreamedChatResponse_SearchQueriesGeneration):
+                    for query in chunk.search_queries:
+                        self.action_events[query.generation_id] = ActionEvent(
+                            action_type="search_query",
+                            init_timestamp=get_ISO_time(),
+                            params=query.text)
+                elif isinstance(chunk, StreamedChatResponse_SearchResults):
                     pass
 
             except Exception as e:
@@ -498,7 +538,7 @@ def override_api(self):
                                 self._override_method(api, method_path, module)
 
                 if api == 'cohere':
-                    # Patch cohere vx.x.x+ methods
+                    # Patch cohere v5.4.0+ methods
                     module_version = version(api)
                     if module_version is None:
                         logger.warning(f'🖇 AgentOps: Cannot determine Cohere version. Only Cohere>=5.4.0 supported.')

examples/cohere_example.py

@@ -0,0 +1,46 @@
+import cohere
+import agentops  # just
+from dotenv import load_dotenv
+load_dotenv()
+
+agentops.init(tags=["cohere", "agentops-demo"])  # three
+co = cohere.Client()
+
+stream = co.chat_stream(
+    message="Tell me everything you can about AgentOps",
+    connectors=[{"id": "web-search"}]
+)
+
+response = ""
+for event in stream:
+    if event.event_type == "text-generation":
+        response += event.text
+        print(event.text, end='')
+    elif event.event_type == "stream-end":
+        print("\n")
+        print(event)
+        print("\n")
+
+stream = co.chat_stream(
+    chat_history=[
+        {"role": "SYSTEM", "message": "You are Adam Silverman: die-hard advocate of AgentOps, leader in AI Agent observability"},
+        {
+            "role": "CHATBOT",
+            "message": "How's your day going? I'd like to tell you about AgentOps: {response}",
+        },
+    ],
+    message="Based on your newfound knowledge of AgentOps, is Cohere a suitable partner for them and how could they integrate?",
+    connectors=[{"id": "web-search"}]
+)
+
+response = ""
+for event in stream:
+    if event.event_type == "text-generation":
+        response += event.text
+        print(event.text, end='')
+    elif event.event_type == "stream-end":
+        print("\n")
+        print(event)
+        print("\n")
+
+agentops.end_session('Success')  # lines