support tracking via lunary (copy of commit from the old repo)

motleycrew/agent/crewai/crewai.py

@@ -1,7 +1,9 @@
-from typing import Any, Optional, Sequence
+from typing import Any, Optional, Sequence, Callable, Dict, List
 
 from crewai import Agent
 from langchain_core.runnables import RunnableConfig
+from langchain.tools.render import render_text_description
+from pydantic import Field
 
 from motleycrew.agent.parent import MotleyAgentAbstractParent
 from motleycrew.agent.shared import MotleyAgentParent
@@ -12,6 +14,59 @@
 from motleycrew.common.utils import to_str
 from motleycrew.common import LLMFramework
 from motleycrew.common.llms import init_llm
+from motleycrew.tracking import add_default_callbacks_to_langchain_config
+
+
+class CrewAIAgentWithConfig(Agent):
+
+    def execute_task(
+            self,
+            task: Any,
+            context: Optional[str] = None,
+            tools: Optional[List[Any]] = None,
+            config: Optional[RunnableConfig] = None
+    ) -> str:
+        """Execute a task with the agent.
+
+        Args:
+            task: Task to execute.
+            context: Context to execute the task in.
+            tools: Tools to use for the task.
+            config: Runnable config
+        Returns:
+            Output of the agent
+        """
+        task_prompt = task.prompt()
+
+        if context:
+            task_prompt = self.i18n.slice("task_with_context").format(
+                task=task_prompt, context=context
+            )
+
+        tools = self._parse_tools(tools or self.tools)
+        self.create_agent_executor(tools=tools)
+        self.agent_executor.tools = tools
+        self.agent_executor.task = task
+        self.agent_executor.tools_description = render_text_description(tools)
+        self.agent_executor.tools_names = self.__tools_names(tools)
+
+        result = self.agent_executor.invoke(
+            {
+                "input": task_prompt,
+                "tool_names": self.agent_executor.tools_names,
+                "tools": self.agent_executor.tools_description,
+            },
+            config=config
+        )["output"]
+
+        if self.max_rpm:
+            self._rpm_controller.stop_rpm_counter()
+
+        return result
+
+    @staticmethod
+    def __tools_names(tools) -> str:
+        return ", ".join([t.name for t in tools])
 
 
 class CrewAIMotleyAgentParent(MotleyAgentParent):
@@ -61,8 +116,10 @@ def invoke(
             raise ValueError(f"`task` must be a string or a Task, not {type(task)}")
 
         langchain_tools = [tool.to_langchain_tool() for tool in self.tools.values()]
+        config = add_default_callbacks_to_langchain_config(config)
+
         out = self.agent.execute_task(
-            task, to_str(task.message_history), tools=langchain_tools
+            task, to_str(task.message_history), tools=langchain_tools, config=config
         )
         task.outputs = [out]
         # TODO: extract message history from agent, attach it to the task
@@ -94,7 +151,7 @@ def from_crewai_params(
 
         def agent_factory(tools: dict[str, MotleyTool]):
             langchain_tools = [t.to_langchain_tool() for t in tools.values()]
-            agent = Agent(
+            agent = CrewAIAgentWithConfig(
                 role=role,
                 goal=goal,
                 backstory=backstory,
@@ -116,7 +173,7 @@ def agent_factory(tools: dict[str, MotleyTool]):
 
     @staticmethod
     def from_agent(
-        agent: Agent,
+        agent: CrewAIAgentWithConfig,
         delegation: bool | Sequence[MotleyAgentAbstractParent] = False,
         tools: Sequence[MotleySupportedTool] | None = None,
         verbose: bool = False,

motleycrew/agent/langchain/langchain.py

@@ -10,6 +10,7 @@
 from motleycrew.tasks import Task
 
 from motleycrew.tool import MotleyTool
+from motleycrew.tracking import add_default_callbacks_to_langchain_config
 from motleycrew.common import MotleySupportedTool
 from motleycrew.common import MotleyAgentFactory
 from motleycrew.common import LLMFramework
@@ -44,6 +45,7 @@ def invoke(
         self.materialize()
         self.agent: AgentExecutor
 
+        config = add_default_callbacks_to_langchain_config(config)
         if isinstance(task, str):
             assert self.crew, "can't create a task outside a crew"
             # TODO: feed in context/task.message_history correctly

motleycrew/agent/llama_index/llama_index_react.py

@@ -1,33 +1,39 @@
 from typing import Sequence
 from llama_index.core.agent import ReActAgent
 from llama_index.core.llms import LLM
+from llama_index.core.callbacks import CallbackManager
 
 from motleycrew.agent.llama_index import LlamaIndexMotleyAgentParent
 from motleycrew.agent.parent import MotleyAgentAbstractParent
 from motleycrew.tool import MotleyTool
 from motleycrew.common import MotleySupportedTool
 from motleycrew.common import LLMFramework
 from motleycrew.common.llms import init_llm
+from motleycrew.tracking import get_default_callbacks_list
 
 
 class ReActLlamaIndexMotleyAgent(LlamaIndexMotleyAgentParent):
-    def __init__(self,
-                 goal: str,
-                 name: str | None = None,
-                 delegation: bool | Sequence[MotleyAgentAbstractParent] = False,
-                 tools: Sequence[MotleySupportedTool] | None = None,
-                 llm: LLM | None = None,
-                 verbose: bool = False):
+    def __init__(
+        self,
+        goal: str,
+        name: str | None = None,
+        delegation: bool | Sequence[MotleyAgentAbstractParent] = False,
+        tools: Sequence[MotleySupportedTool] | None = None,
+        llm: LLM | None = None,
+        verbose: bool = False,
+    ):
         if llm is None:
             llm = init_llm(llm_framework=LLMFramework.LLAMA_INDEX)
 
         def agent_factory(tools: dict[str, MotleyTool]):
             llama_index_tools = [t.to_llama_index_tool() for t in tools.values()]
             # TODO: feed goal into the agent's prompt
+            callbacks = get_default_callbacks_list(LLMFramework.LLAMA_INDEX)
             agent = ReActAgent.from_tools(
                 tools=llama_index_tools,
                 llm=llm,
                 verbose=verbose,
+                callback_manager=CallbackManager(callbacks),
             )
             return agent
 
@@ -37,5 +43,5 @@ def agent_factory(tools: dict[str, MotleyTool]):
             agent_factory=agent_factory,
             delegation=delegation,
             tools=tools,
-            verbose=verbose
+            verbose=verbose,
         )

motleycrew/common/enums.py

@@ -5,3 +5,17 @@ class LLMFamily:
 class LLMFramework:
     LANGCHAIN = "langchain"
     LLAMA_INDEX = "llama_index"
+
+
+class LunaryRunType:
+    LLM = "llm"
+    AGENT = "agent"
+    TOOL = "tool"
+    CHAIN = "chain"
+    EMBED = "embed"
+
+
+class LunaryEventName:
+    START = "start"
+    END = "end"
+    ERROR = "error"

motleycrew/tracking/__init__.py

@@ -0,0 +1 @@
+from .utils import add_default_callbacks_to_langchain_config, get_default_callbacks_list