From ccfb6fe84e07907dc859789686cf871009ca552a Mon Sep 17 00:00:00 2001
From: Gyuri Park <inglifestora@naver.com>
Date: Tue, 30 Sep 2025 09:00:05 +0000
Subject: [PATCH] Integrate feature group creation tool into LLM agent

- agent_schema.py
Define Pydantic schema for feature group validation.
- agent_service.py
Implement the create_feature_group tool and integrate it with the DSPy agent.

Issue-Id: AIMLFW-269
Change-Id: If4797dbdc2cde69ee5895a1d9478ade22ae1f359
Signed-off-by: Gyuri Park <inglifestora@naver.com>
---
 trainingmgr/schemas/agent_schema.py  | 34 +++++++++++++++++++
 trainingmgr/service/agent_service.py | 66 ++++++++++++++++++++++++++++++++----
 2 files changed, 93 insertions(+), 7 deletions(-)
 create mode 100644 trainingmgr/schemas/agent_schema.py

diff --git a/trainingmgr/schemas/agent_schema.py b/trainingmgr/schemas/agent_schema.py
new file mode 100644
index 0000000..21d7702
--- /dev/null
+++ b/trainingmgr/schemas/agent_schema.py
@@ -0,0 +1,34 @@
+# ==================================================================================
+#
+#       Copyright (c) 2025 Gyuri Park <inglifestora@naver.com> All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#          http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+# ==================================================================================
+from pydantic import BaseModel
+from typing import Optional
+
+class FeatureGroupIntent(BaseModel):
+    featuregroup_name: str
+    feature_list: str
+    datalake_source: str
+    enable_dme: bool
+    host: str
+    port: str
+    bucket: str
+    token: str
+    measurement: str
+    db_org: str
+    dme_port: Optional[str] = None
+    source_name: Optional[str] = None
+    measured_obj_class: Optional[str] = None
\ No newline at end of file
diff --git a/trainingmgr/service/agent_service.py b/trainingmgr/service/agent_service.py
index d1aacc8..c1a76e7 100644
--- a/trainingmgr/service/agent_service.py
+++ b/trainingmgr/service/agent_service.py
@@ -15,13 +15,66 @@
 #   limitations under the License.
 #
 # ==================================================================================
-import dspy
 import os
+import requests
+import dspy
+from threading import Lock
 from trainingmgr.common.trainingmgr_config import TrainingMgrConfig
 from trainingmgr.common.exceptions_utls import TMException
-from threading import Lock
+from trainingmgr.schemas.agent_schema import FeatureGroupIntent
+
+CONFIG = TrainingMgrConfig()
+LOGGER = CONFIG.logger
+
+# Define the DSPy tool
+@dspy.Tool
+def create_feature_group(
+    featuregroup_name: str,
+    feature_list: str,
+    enable_dme: bool,
+    host: str,
+    port: str,
+    bucket: str,
+    token: str,
+    measurement: str,
+    db_org: str,
+    dme_port: str = None,
+    source_name: str = None,
+    measured_obj_class: str = None,
+    datalake_source: str = "InfluxSource"
+) -> str:
+    """Create a feature group using the Training Manager API."""
+    try:
+        data = {
+            "featuregroup_name": featuregroup_name,
+            "feature_list": feature_list,
+            "datalake_source": datalake_source,
+            "enable_dme": enable_dme,
+            "host": host,
+            "port": port,
+            "bucket": bucket,
+            "token": token,
+            "measurement": measurement,
+            "db_org": db_org,
+            "dme_port": dme_port,
+            "source_name": source_name,
+            "measured_obj_class": measured_obj_class,
+        }
+        obj = FeatureGroupIntent.model_validate(data)
+        json_payload = obj.model_dump(exclude_none=True)
+
+        tm_ip = CONFIG.my_ip
+        tm_port = CONFIG.my_port
+        if not tm_ip or not tm_port:
+            raise TMException("Training manager IP/Port not configured")
+
+        url = f"http://{tm_ip}:{tm_port}/ai-ml-model-training/v1/featureGroup"
+        response = requests.post(url, json=json_payload, timeout=15)
+        response.raise_for_status()
+        return f"Feature group '{obj.featuregroup_name}' created (status={response.status_code})."
+    except Exception as err:
+        raise TMException(f"Error creating feature group: {str(err)}")
 
-LOGGER = TrainingMgrConfig().logger
 
 # Define the agent signature
 class AgentSignature(dspy.Signature):
@@ -71,7 +124,7 @@ class AgentClient:
             # Agent configuration
             self._agent = dspy.ReAct(
                 AgentSignature,
-                tools=[],
+                tools=[create_feature_group],
                 max_iters=6
             )
 
@@ -80,7 +133,7 @@ class AgentClient:
             return True
 
         except Exception as err:
-            raise TMException(f"fail to initialize agent exception : {str(err)}")
+            raise TMException(f"fail to initialize agent exception: {str(err)}")
 
     def process_user_request(self, user_text_request):
         """Process user request with agent tools."""
@@ -99,5 +152,4 @@ class AgentClient:
             return {
                 'success': False,
                 'error': str(err),
-            }
-
+            }
\ No newline at end of file
-- 
2.16.6