From 2a615fbff38da9cc98c2c43dbbc7524a3f2f206e Mon Sep 17 00:00:00 2001
From: Pratyush Shukla <ps4534@nyu.edu>
Date: Mon, 18 Nov 2024 18:13:58 +0530
Subject: [PATCH] init

---
 taskweaver/llm/__init__.py | 22 ++++++++++++++++++++++
 taskweaver/llm/base.py     | 15 +++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/taskweaver/llm/__init__.py b/taskweaver/llm/__init__.py
index 967ec617..02dda5a0 100644
--- a/taskweaver/llm/__init__.py
+++ b/taskweaver/llm/__init__.py
@@ -11,6 +11,7 @@
     ExtLLMModuleConfig,
     LLMModuleConfig,
     LLMServiceConfig,
+    AgentOpsModuleConfig,
 )
 from taskweaver.llm.google_genai import GoogleGenAIService
 from taskweaver.llm.groq import GroqService, GroqServiceConfig
@@ -48,11 +49,13 @@ def __init__(
         config: LLMModuleConfig,
         injector: Injector,
         ext_llms_config: Optional[ExtLLMModuleConfig] = None,
+        agentops_config: Optional[AgentOpsModuleConfig] = None,
     ):
         self.config = config
         self.injector = injector
         self.ext_llm_injector = Injector([])
         self.ext_llms = {}  # extra llm models
+        self.agentops_config = agentops_config
 
         if self.config.api_type in ["openai", "azure", "azure_ad"]:
             self._set_completion_service(OpenAIService)
@@ -115,6 +118,25 @@ def __init__(
                 llm_completion_service = self._get_completion_service(config)
                 self.ext_llms[key] = llm_completion_service
 
+        if self.agentops_config.enabled:
+            try:
+                from agentops import Client
+                from agentops.llms import LlmTracker
+
+                if Client().is_initialized:
+                    self.agentops_client = Client().current_client
+                else:
+                    self.agentops_client = Client()
+                    self.agentops_client.configure(api_key=self.agentops_config.api_key, default_tags=self.agentops_config.default_tags)
+
+                self.llm_tracker = LlmTracker(self.agentops_client)
+                self.llm_tracker.override_api()
+            except ImportError as exc:
+                raise ImportError(
+                    "AgentOps is not installed. Please install it with 'pip install agentops' "
+                    "to use AgentOps tracking features."
+                ) from exc
+
     def _set_completion_service(self, svc: Type[CompletionService]) -> None:
         self.completion_service: CompletionService = self.injector.get(svc)
         self.injector.binder.bind(svc, to=self.completion_service)
diff --git a/taskweaver/llm/base.py b/taskweaver/llm/base.py
index f29d943c..c4e612f4 100644
--- a/taskweaver/llm/base.py
+++ b/taskweaver/llm/base.py
@@ -76,6 +76,21 @@ def _set_name(self, name: str) -> None:
         self.name = f"llm.{name}"
 
 
+class AgentOpsModuleConfig(ModuleConfig):
+    def _configure(self) -> None:
+        self._set_name("agentops")
+        self.enabled: bool = self._get_bool("enabled", False)
+        self.api_key: Optional[str] = self._get_str(
+            "api_key",
+            None,
+            required=False,
+        )
+        self.default_tags: Optional[List[str]] = self._get_list(
+            "default_tags",
+            ["taskweaver"],
+        )
+
+
 class CompletionService(abc.ABC):
     @abc.abstractmethod
     def chat_completion(