From 8f23c29818daacfc1be8e41aadb5588dcf70c64d Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 16 Nov 2023 09:44:17 -0300
Subject: [PATCH] Refactor StoreService upload method to return
 CreateComponentResponse

---
 src/backend/langflow/services/store/service.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/backend/langflow/services/store/service.py b/src/backend/langflow/services/store/service.py
index 6833d84ed..9257f0a21 100644
--- a/src/backend/langflow/services/store/service.py
+++ b/src/backend/langflow/services/store/service.py
@@ -1,18 +1,19 @@
 import json
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 from uuid import UUID
 
 import httpx
 from httpx import HTTPError, HTTPStatusError
+from loguru import logger
+
 from langflow.services.base import Service
 from langflow.services.store.schema import (
-    ComponentResponse,
+    CreateComponentResponse,
     DownloadComponentResponse,
     ListComponentResponse,
     StoreComponentCreate,
 )
 from langflow.services.store.utils import process_tags_for_post
-from loguru import logger
 
 if TYPE_CHECKING:
     from langflow.services.settings.service import SettingsService
@@ -74,7 +75,7 @@ class StoreService(Service):
 
     async def _get(
         self, url: str, api_key: Optional[str] = None, params: Optional[Dict[str, Any]] = None
-    ) -> List[Dict[str, Any]]:
+    ) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
         """Utility method to perform GET requests."""
         if api_key:
             headers = {"Authorization": f"Bearer {api_key}"}
@@ -257,7 +258,7 @@ class StoreService(Service):
 
         return DownloadComponentResponse(**component)
 
-    async def upload(self, api_key: str, component_data: StoreComponentCreate) -> ComponentResponse:
+    async def upload(self, api_key: str, component_data: StoreComponentCreate) -> CreateComponentResponse:
         headers = {"Authorization": f"Bearer {api_key}"}
         component_dict = component_data.dict(exclude_unset=True)
         # Parent is a UUID, but the store expects a string
@@ -273,7 +274,7 @@ class StoreService(Service):
                 response = await client.post(self.components_url, headers=headers, json=component_dict)
                 response.raise_for_status()
             component = response.json()["data"]
-            return ComponentResponse(**component)
+            return CreateComponentResponse(**component)
         except HTTPError as exc:
             if response:
                 try: