[FEAT] TaskWeaver Integration

[the-praxs]

📥 Pull Request

📘 Description

• Closes [Feature]: TaskWeaver Support #500
• Adds support for TaskWeaver by using the SessionEventHandler class
• Adds support for tracking the User, Planner, and CodeInterpreter agents present in TaskWeaver

🧪 Testing
Tested using the below code by checking the events in the dashboard -

import agentops
from dotenv import load_dotenv
import os

from taskweaver.app.app import TaskWeaverApp
from agentops.partners.taskweaver_event_handler import TaskWeaverEventHandler

load_dotenv()
AGENTOPS_API_KEY = os.getenv("AGENTOPS_API_KEY")

agentops.init(api_key=AGENTOPS_API_KEY, default_tags=["taskweaver", "test"])

# Assuming the current project directory contains the taskweaver_config.json file
app_dir = "./project/"
app = TaskWeaverApp(app_dir=app_dir)
session = app.get_session()
handler = TaskWeaverEventHandler()

session.event_emitter.register(handler)

user_query = "hello, what can you do?"
# response_round = session.send_message(user_query, event_handler=handler)
response_round = session.send_message(user_query)
print(response_round.to_dict())

⚠️ Caveats

• Every information that we need to track is not present out of the box in the SessionEventHandler class. This is a TaskWeaver architecture-related problem that we need to resolve.
• We assume for brevity that the default agent is Planner since it is the intermediary between the User and the CodeInterpreter. This is not a scalable solution and can be the root cause of problems in the future.
• Extensive testing is required with different and preferably complex codes using TaskWeaver in library mode.

🗒️ To Do

• Resolve above mentioned caveats.
• Create docs page linking to the example.
• Update the README to include a link to the docs page.

[codecov]

Codecov Report

Attention: Patch coverage is 10.71429% with 150 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
agentops/llms/providers/taskweaver.py	18.27%	76 Missing ⚠️
agentops/partners/taskweaver_event_handler.py	0.00%	67 Missing ⚠️
agentops/llms/tracker.py	12.50%	7 Missing ⚠️

Flag	Coverage Δ
unittests	52.35% <10.71%> (-1.66%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
agentops/llms/tracker.py	23.36% <12.50%> (-0.88%)	⬇️
agentops/partners/taskweaver_event_handler.py	0.00% <0.00%> (ø)
agentops/llms/providers/taskweaver.py	18.27% <18.27%> (ø)

[teocns]

Every information that we need to track is not present out of the box in the SessionEventHandler class. This is a TaskWeaver architecture-related problem that we need to resolve.

Can you share further findings i.e the data set / schema yielded by TaskWeaver's?

[areibman]

Also todo:

• Create docs page linking to the example
• Update the README to include a link to the docs page

[the-praxs]

Currently I am working on tracking the LLM information for the LLMEvent object.

For this I need to have the chat_completion and chat_completion_stream methods in LLMApi class exposed within the TaskWeaver class object. I am working on a quick hack for this.

[teocns]

For this I need to have the chat_completion and chat_completion_stream methods in LLMApi class exposed within the TaskWeaver class object. I am working on a quick hack for this.

hey could you share backlinks and reference to the code to make it easy to jump in? I want to have a second hand pair of eyes on this to make this move faster

[the-praxs]

For this I need to have the chat_completion and chat_completion_stream methods in LLMApi class exposed within the TaskWeaver class object. I am working on a quick hack for this.

hey could you share backlinks and reference to the code to make it easy to jump in? I want to have a second hand pair of eyes on this to make this move faster

TaskWeaver runs as a CLI application in isolation mode by default but you can use it as a library too. TaskWeaver extensively uses the Injector library (API reference here]) in their codebase to inject the dependencies.

This is the examples code from the docs that we are also using to test the integration -

from taskweaver.app.app import TaskWeaverApp

# This is the folder that contains the taskweaver_config.json file and not the repo root. Defaults to "./project/"
app_dir = "./project/"
app = TaskWeaverApp(app_dir=app_dir)
session = app.get_session()

user_query = "hello, what can you do?"
response_round = session.send_message(user_query)
print(response_round.to_dict())

The app_dir is the project directory which will contain all the code and logs generated by TaskWeaver. All the configurations are defined in a file called task weaver_config.json which will contain the configurations for different items based on the classes defined in taskweaver/llm/base.py

The TaskWeaverApp class contains the app_injector attribute which is the Injector object to bind and get the dependencies. In our case, we have the LlmApi class which is injected into the TaskWeaverApp object and is used by Planner and CodeGenerator classes.

In LlmApi you are setting up the CompletionService for the different LLM providers (OpenAI, Anthropic, Qwen etc.) and injecting them using the injector object. Based on the configurations from task weaver_config.json, the chat_completion of the configured LLM(s) is injected via the _set_completion_service method and their chat_completion method is set to the chat_completion method of the LlmApi class.

This method is now used to chat with the LLM and output the responses. This is what we need to monkey patch. We do not have a direct access to the LlmApi instance in the TaskWeaverApp object but we are exposed to the app_injector object which is of type Injector. Using this we can use the code below to fetch the configured LlmApi class -

from taskweaver.llm import LLMApi
llmapi = app.app_injector.get(LLMApi)

Now we have the chat_completion method exposed for us which can be monkey patched. To bind it back we can use the app_injector object again. For example, if we were to substitute OpenAIService with AnthropicService we would use the following code -

from taskweaver.llm.anthropic import AnthropicService

llmapi.completion_service = llmapi.injector.get(AnthropicService)
llmapi.injector.binder.bind(AnthropicService, to=llmapi.completion_service)

app.app_injector.binder.bind(LLMApi, to=llmapi)

Also - to override the chat_completion function directly, we can use the following code -

from taskweaver.llm.anthropic import AnthropicService

llmapi.chat_completion = AnthropicService.chat_completion
app.app_injector.binder.bind(LLMApi, to=llmapi)

[areibman]

Awesome!!! Almost at the finish line. Can we do 2 things:

1. When the AgentOps session is created, can you please insert "TaskWeaver" as one of the tags?
2. Can you add a short documentation page for how to run it?

[areibman]

Very nice-- elegant!