From ca744481ee0355a66e37e6322be6bffc44391ebc Mon Sep 17 00:00:00 2001
From: Nikita Ruban <nikitkaruban@gmail.com>
Date: Thu, 27 Mar 2025 15:28:03 +0000
Subject: [PATCH] fix: Remove is enabled check for raw data requests

- Leave enabled check only for llm agent usage
- Raw data tool will be enabled if a corresponding api key has been provided in env
---
 src/__tests__/SentientAI.test.ts | 17 ++++-------------
 src/sentientAI.ts                |  7 ++-----
 2 files changed, 6 insertions(+), 18 deletions(-)

diff --git a/src/__tests__/SentientAI.test.ts b/src/__tests__/SentientAI.test.ts
index 473e459..7ca3426 100644
--- a/src/__tests__/SentientAI.test.ts
+++ b/src/__tests__/SentientAI.test.ts
@@ -60,14 +60,14 @@ describe("SentientAI", () => {
   it("should throw if FAST_LLM_MODEL is not set", () => {
     delete process.env.FAST_LLM_MODEL;
     expect(() => new SentientAI()).toThrow(
-      "FAST_LLM_MODEL and LLM_MODEL must be set",
+      "FAST_LLM_MODEL and LLM_MODEL must be set"
     );
   });
 
   it("should throw if LLM_MODEL is not set", () => {
     delete process.env.LLM_MODEL;
     expect(() => new SentientAI()).toThrow(
-      "FAST_LLM_MODEL and LLM_MODEL must be set",
+      "FAST_LLM_MODEL and LLM_MODEL must be set"
     );
   });
 
@@ -75,7 +75,7 @@ describe("SentientAI", () => {
     delete process.env.FAST_LLM_MODEL;
     delete process.env.LLM_MODEL;
     expect(() => new SentientAI()).toThrow(
-      "FAST_LLM_MODEL and LLM_MODEL must be set",
+      "FAST_LLM_MODEL and LLM_MODEL must be set"
     );
   });
 
@@ -88,7 +88,6 @@ describe("SentientAI", () => {
 
   describe("getRawData", () => {
     it("should return raw data for an enabled tool", async () => {
-      process.env.ENABLED_TOOLS = "weather-current";
       const sentai = new SentientAI();
       const params = { latitude: 37.7749, longitude: -122.4194 };
       const result = await sentai.getRawData("weather-current", params);
@@ -99,15 +98,7 @@ describe("SentientAI", () => {
       process.env.ENABLED_TOOLS = "weather-current";
       const sentai = new SentientAI();
       await expect(sentai.getRawData("non-existent-tool", {})).rejects.toThrow(
-        "Tool 'non-existent-tool' not found",
-      );
-    });
-
-    it("should throw when tool is not enabled", async () => {
-      process.env.ENABLED_TOOLS = "weather-current";
-      const sentai = new SentientAI();
-      await expect(sentai.getRawData("weather-forecast", {})).rejects.toThrow(
-        "Tool 'weather-forecast' is not enabled",
+        "Tool 'non-existent-tool' not found"
       );
     });
   });
diff --git a/src/sentientAI.ts b/src/sentientAI.ts
index 5a61a58..63a91d0 100644
--- a/src/sentientAI.ts
+++ b/src/sentientAI.ts
@@ -36,7 +36,7 @@ export class SentientAI {
     toolName: string,
     params: Record<string, any>
   ): Promise<any> {
-    const tool = this.getTool(toolName);
+    const tool = this.getRawTool(toolName);
     return this.rawDataProvider.process(tool, params);
   }
 
@@ -48,14 +48,11 @@ export class SentientAI {
     return this.orchestrator.processStream(input);
   }
 
-  private getTool(toolName: string): QSTool {
+  private getRawTool(toolName: string): QSTool {
     const tool = ToolRegistry.getTool(toolName as ToolName);
     if (!tool) {
       throw new Error(`Tool '${toolName}' not found`);
     }
-    if (!ToolRegistry.isEnabled(toolName as ToolName)) {
-      throw new Error(`Tool '${toolName}' is not enabled`);
-    }
     return tool;
   }
 }