[FEAT] TaskWeaver Integration

agentops/llms/providers/taskweaver.py

@@ -0,0 +1,146 @@
+import pprint
+from typing import Optional
+import json
+
+from agentops.event import ErrorEvent, LLMEvent, ActionEvent
+from agentops.session import Session
+from agentops.log_config import logger
+from agentops.helpers import get_ISO_time, check_call_stack_for_agent_id
+from agentops.llms.providers.instrumented_provider import InstrumentedProvider
+from agentops.singleton import singleton
+
+
+@singleton
+class TaskWeaverProvider(InstrumentedProvider):
+    original_chat_completion = None
+
+    def __init__(self, client):
+        super().__init__(client)
+        self._provider_name = "TaskWeaver"
+        self.client.add_default_tags(["taskweaver"])
+
+    def handle_response(self, response, kwargs, init_timestamp, session: Optional[Session] = None) -> dict:
+        """Handle responses for TaskWeaver"""
+        llm_event = LLMEvent(init_timestamp=init_timestamp, params=kwargs)
+        action_event = ActionEvent(init_timestamp=init_timestamp)
+
+        try:
+            response_dict = response.get("response", {})
+
+            action_event.params = kwargs.get("json_schema", None)
+            action_event.returns = response_dict
+            action_event.end_timestamp = get_ISO_time()
+            self._safe_record(session, action_event)
+        except Exception as e:
+            error_event = ErrorEvent(
+                trigger_event=action_event, exception=e, details={"response": str(response), "kwargs": str(kwargs)}
+            )
+            self._safe_record(session, error_event)
+            kwargs_str = pprint.pformat(kwargs)
+            response_str = pprint.pformat(response)
+            logger.error(
+                f"Unable to parse response for Action call. Skipping upload to AgentOps\n"
+                f"response:\n {response_str}\n"
+                f"kwargs:\n {kwargs_str}\n"
+            )
+
+        try:
+            llm_event.init_timestamp = init_timestamp
+            llm_event.params = kwargs
+            llm_event.returns = response_dict
+            llm_event.agent_id = check_call_stack_for_agent_id()
+            llm_event.model = kwargs.get("model", "unknown")
+            llm_event.prompt = kwargs.get("messages")
+            llm_event.completion = response_dict.get("message", "")
+            llm_event.end_timestamp = get_ISO_time()
+            self._safe_record(session, llm_event)
+
+        except Exception as e:
+            error_event = ErrorEvent(
+                trigger_event=llm_event, exception=e, details={"response": str(response), "kwargs": str(kwargs)}
+            )
+            self._safe_record(session, error_event)
+            kwargs_str = pprint.pformat(kwargs)
+            response_str = pprint.pformat(response)
+            logger.error(
+                f"Unable to parse response for LLM call. Skipping upload to AgentOps\n"
+                f"response:\n {response_str}\n"
+                f"kwargs:\n {kwargs_str}\n"
+            )
+
+        return response
+
+    def override(self):
+        """Override TaskWeaver's chat completion methods"""
+        try:
+            from taskweaver.llm import llm_completion_config_map
+
+            def create_patched_chat_completion(original_method):
+                """Create a new patched chat_completion function with bound original method"""
+
+                def patched_chat_completion(service, *args, **kwargs):
+                    init_timestamp = get_ISO_time()
+                    session = kwargs.get("session", None)
+                    if "session" in kwargs.keys():
+                        del kwargs["session"]
+
+                    result = original_method(service, *args, **kwargs)
+                    kwargs.update(
+                        {
+                            "model": self._get_model_name(service),
+                            "messages": args[0],
+                            "stream": args[1],
+                            "temperature": args[2],
+                            "max_tokens": args[3],
+                            "top_p": args[4],
+                            "stop": args[5],
+                        }
+                    )
+
+                    if kwargs["stream"]:
+                        accumulated_content = ""
+                        for chunk in result:
+                            if isinstance(chunk, dict) and "content" in chunk:
+                                accumulated_content += chunk["content"]
+                            else:
+                                accumulated_content += chunk
+                            yield chunk
+                        accumulated_content = json.loads(accumulated_content)
+                        return self.handle_response(accumulated_content, kwargs, init_timestamp, session=session)
+                    else:
+                        return self.handle_response(result, kwargs, init_timestamp, session=session)
+
+                return patched_chat_completion
+
+            for service_name, service_class in llm_completion_config_map.items():
+                if not hasattr(service_class, "_original_chat_completion"):
+                    service_class._original_chat_completion = service_class.chat_completion
+                    service_class.chat_completion = create_patched_chat_completion(
+                        service_class._original_chat_completion
+                    )
+
+        except Exception as e:
+            logger.error(f"Failed to patch method: {str(e)}", exc_info=True)
+
+    def undo_override(self):
+        """Restore original TaskWeaver chat completion methods"""
+        try:
+            from taskweaver.llm import llm_completion_config_map
+
+            for service_name, service_class in llm_completion_config_map.items():
+                service_class.chat_completion = service_class._original_chat_completion
+                delattr(service_class, "_original_chat_completion")
+
+        except Exception as e:
+            logger.error(f"Failed to restore original method: {str(e)}", exc_info=True)
+
+    def _get_model_name(self, service) -> str:
+        """Extract model name from service instance"""
+        model_name = "unknown"
+        if hasattr(service, "config"):
+            config = service.config
+            if hasattr(config, "model"):
+                model_name = config.model or "unknown"
+            elif hasattr(config, "llm_module_config") and hasattr(config.llm_module_config, "model"):
+                model_name = config.llm_module_config.model or "unknown"
+        return model_name

agentops/llms/tracker.py

@@ -16,6 +16,7 @@
 from .providers.anthropic import AnthropicProvider
 from .providers.mistral import MistralProvider
 from .providers.ai21 import AI21Provider
+from .providers.taskweaver import TaskWeaverProvider
 
 original_func = {}
 original_create = None
@@ -54,6 +55,9 @@
                 "client.answer.create",
             ),
         },
+        "taskweaver": {
+            "0.0.1": ("chat_completion", "chat_completion_stream"),
+        },
     }
 
     def __init__(self, client):
@@ -164,6 +168,15 @@
                     else:
                         logger.warning(f"Only LlamaStackClient>=0.0.53 supported. v{module_version} found.")
 
+                if api == "taskweaver":
+                    module_version = version(api)
+
+                    if Version(module_version) >= parse("0.0.1"):
+                        provider = TaskWeaverProvider(self.client)
+                        provider.override()
+                    else:
+                        logger.warning(f"Only TaskWeaver>=0.0.1 supported. v{module_version} found.")
+
     def stop_instrumenting(self):
         OpenAiProvider(self.client).undo_override()
         GroqProvider(self.client).undo_override()
@@ -174,3 +187,4 @@
         MistralProvider(self.client).undo_override()
         AI21Provider(self.client).undo_override()
         LlamaStackClientProvider(self.client).undo_override()
+        TaskWeaverProvider(self.client).undo_override()

agentops/partners/taskweaver_event_handler.py

@@ -0,0 +1,156 @@
+from taskweaver.module.event_emitter import (
+    SessionEventHandlerBase,
+    TaskWeaverEvent,
+    EventScope,
+    SessionEventType,
+    RoundEventType,
+    PostEventType,
+)
+import agentops
+from typing import Dict, Optional, Any
+from uuid import UUID
+
+
+class TaskWeaverEventHandler(SessionEventHandlerBase):
+    def __init__(self):
+        super().__init__()
+        self.current_round_id: Optional[str] = None
+        self.agent_sessions: Dict[str, Any] = {}
+        self._message_buffer: Dict[str, str] = {}
+        self._attachment_buffer: Dict[str, Dict[str, Any]] = {}
+        self._active_agents: Dict[str, str] = {}  # Maps role_round_id to agent_id
+
+    def _get_or_create_agent(self, role: str) -> str:
+        """Get existing agent ID or create new agent for role+round combination"""
+        agent_key = f"{role}"
+        if agent_key not in self._active_agents:
+            agent_id = agentops.create_agent(name=role)
+            if agent_id:  # Only store if agent creation was successful
+                self._active_agents[agent_key] = agent_id
+        return self._active_agents.get(agent_key)
+
+    def handle_session(self, type: SessionEventType, msg: str, extra: Any, **kwargs: Any):
+        agentops.record(
+            agentops.ActionEvent(action_type=type.value, params={"message": msg}, returns=str(extra) if extra else None)
+        )
+
+    def handle_round(self, type: RoundEventType, msg: str, extra: Any, round_id: str, **kwargs: Any):
+        if type == RoundEventType.round_start:
+            self.current_round_id = round_id
+            agentops.record(agentops.ActionEvent(action_type="round_start", params={"round_id": round_id}, returns=msg))
+
+        elif type == RoundEventType.round_error:
+            agentops.record(
+                agentops.ErrorEvent(error_type="round_error", details={"message": msg, "round_id": round_id})
+            )
+
+        elif type == RoundEventType.round_end:
+            agentops.record(agentops.ActionEvent(action_type="round_end", params={"round_id": round_id}, returns=msg))
+            self.current_round_id = None
+
+    def handle_post(self, type: PostEventType, msg: str, extra: Any, post_id: str, round_id: str, **kwargs: Any):
+        role = extra.get("role", "Planner")
+        agent_id = self._get_or_create_agent(role=role)
+
+        if type == PostEventType.post_start:
+            agentops.record(
+                agentops.ActionEvent(
+                    action_type="post_start",
+                    params={
+                        "post_id": post_id,
+                        "round_id": round_id,
+                        "agent_id": agent_id,
+                    },
+                    returns=msg,
+                )
+            )
+
+        elif type == PostEventType.post_message_update:
+            is_end = extra.get("is_end", False)
+            if not is_end:
+                self._message_buffer[post_id] = self._message_buffer.get(post_id, "") + msg
+            else:
+                agentops.record(
+                    agentops.ActionEvent(
+                        action_type="post_message_update",
+                        params={
+                            "post_id": post_id,
+                            "round_id": round_id,
+                            "agent_id": agent_id,
+                            "is_end": is_end,
+                            "model": extra.get("model", None),
+                        },
+                        returns=self._message_buffer.get(post_id, ""),
+                    )
+                )
+
+            if is_end:
+                self._message_buffer.pop(post_id, None)
+
+        elif type == PostEventType.post_attachment_update:
+            attachment_id = extra.get("id", "")
+            attachment_type = extra.get("type", "")
+            is_end = extra.get("is_end", False)
+
+            if attachment_id not in self._attachment_buffer:
+                self._attachment_buffer[attachment_id] = {
+                    "type": attachment_type,
+                    "content": "",
+                    "post_id": post_id,
+                    "round_id": round_id,
+                    "agent_id": agent_id,
+                }
+
+                agentops.record(
+                    agentops.ActionEvent(
+                        action_type="attachment_stream_start",
+                        params={
+                            "attachment_id": attachment_id,
+                            "attachment_type": str(attachment_type),
+                            "post_id": post_id,
+                            "round_id": round_id,
+                            "agent_id": agent_id,
+                        },
+                    )
+                )
+
+            self._attachment_buffer[attachment_id]["content"] += str(msg)
+
+            if is_end:
+                buffer = self._attachment_buffer[attachment_id]
+                agentops.record(
+                    agentops.ToolEvent(
+                        name=str(buffer["type"]),
+                        params={
+                            "post_id": buffer["post_id"],
+                            "round_id": buffer["round_id"],
+                            "attachment_id": attachment_id,
+                        },
+                        returns=buffer["content"],
+                        agent_id=buffer["agent_id"],
+                    )
+                )
+                self._attachment_buffer.pop(attachment_id)
+
+        elif type == PostEventType.post_error:
+            agentops.record(
+                agentops.ErrorEvent(
+                    error_type="post_error",
+                    details={"message": msg, "post_id": post_id, "round_id": round_id},
+                )
+            )
+
+        elif type == PostEventType.post_end:
+            agentops.record(
+                agentops.ActionEvent(
+                    action_type="post_end",
+                    params={"post_id": post_id, "round_id": round_id, "agent_id": agent_id},
+                    returns=msg,
+                )
+            )
+
+    def cleanup_round(self, round_id: str):
+        """Cleanup agents and buffers for a completed round"""
+        self._active_agents = {k: v for k, v in self._active_agents.items() if not k.endswith(round_id)}
+        self._message_buffer.clear()
+        self._attachment_buffer.clear()