From 526ea9044e69892ff418ba32f3efae64535ef7d8 Mon Sep 17 00:00:00 2001
From: Yonatan Shelach <92271540+yonishelach@users.noreply.github.com>
Date: Mon, 4 Nov 2024 15:30:41 +0200
Subject: [PATCH] [Controller] Move `_process_output` to SQL Client (#35)

* move function to child because uses unmandatory method

* format
---
 controller/src/controller/db/client.py        | 24 -------------------
 controller/src/controller/db/sql/sqlclient.py | 24 +++++++++++++++++++
 genai_factory/src/genai_factory/__main__.py   |  2 +-
 .../workflows/workflow_server.py              |  1 +
 4 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/controller/src/controller/db/client.py b/controller/src/controller/db/client.py
index f19875f..a133116 100644
--- a/controller/src/controller/db/client.py
+++ b/controller/src/controller/db/client.py
@@ -706,27 +706,3 @@ def list_sessions(
         :return: The list of chat sessions.
         """
         pass
-
-    def _process_output(
-        self,
-        items,
-        obj_class,
-        mode: api_models.OutputMode = api_models.OutputMode.DETAILS,
-    ) -> Union[list, dict]:
-        """
-        Process the output of a query. Use this method to convert the output to the desired format.
-        For example when listing.
-
-        :param items:     The items to process.
-        :param obj_class: The class of the items.
-        :param mode:      The output mode.
-
-        :return: The processed items.
-        """
-        if mode == api_models.OutputMode.NAMES:
-            return [item.name for item in items]
-        items = [self._from_db_object(item, obj_class) for item in items]
-        if mode == api_models.OutputMode.DETAILS:
-            return items
-        short = mode == api_models.OutputMode.SHORT
-        return [item.to_dict(short=short) for item in items]
diff --git a/controller/src/controller/db/sql/sqlclient.py b/controller/src/controller/db/sql/sqlclient.py
index 7ed0b00..7c52610 100644
--- a/controller/src/controller/db/sql/sqlclient.py
+++ b/controller/src/controller/db/sql/sqlclient.py
@@ -1360,3 +1360,27 @@ def list_sessions(
         if last > 0:
             query = query.limit(last)
         return self._process_output(query.all(), api_models.ChatSession, output_mode)
+
+    def _process_output(
+        self,
+        items,
+        obj_class,
+        mode: api_models.OutputMode = api_models.OutputMode.DETAILS,
+    ) -> Union[list, dict]:
+        """
+        Process the output of a query. Use this method to convert the output to the desired format.
+        For example when listing.
+
+        :param items:     The items to process.
+        :param obj_class: The class of the items.
+        :param mode:      The output mode.
+
+        :return: The processed items.
+        """
+        if mode == api_models.OutputMode.NAMES:
+            return [item.name for item in items]
+        items = [self._to_schema_object(item, obj_class) for item in items]
+        if mode == api_models.OutputMode.DETAILS:
+            return items
+        short = mode == api_models.OutputMode.SHORT
+        return [item.to_dict(short=short) for item in items]
diff --git a/genai_factory/src/genai_factory/__main__.py b/genai_factory/src/genai_factory/__main__.py
index 9a4900d..608bdc9 100644
--- a/genai_factory/src/genai_factory/__main__.py
+++ b/genai_factory/src/genai_factory/__main__.py
@@ -20,9 +20,9 @@
 
 import click
 import dotenv
-from genai_factory.api import router
 
 from genai_factory import WorkflowServerConfig
+from genai_factory.api import router
 
 # Load the environment variables:
 dotenv.load_dotenv(os.environ.get("GENAI_FACTORY_ENV_PATH", "./.env"))
diff --git a/genai_factory/src/genai_factory/workflows/workflow_server.py b/genai_factory/src/genai_factory/workflows/workflow_server.py
index 4cfe865..2c04fe8 100644
--- a/genai_factory/src/genai_factory/workflows/workflow_server.py
+++ b/genai_factory/src/genai_factory/workflows/workflow_server.py
@@ -15,6 +15,7 @@
 from urllib.parse import urlparse
 
 import uvicorn
+
 from genai_factory.config import WorkflowServerConfig
 from genai_factory.controller_client import ControllerClient
 from genai_factory.schemas import WorkflowType