Refactor StoreService upload method to return

CreateComponentResponse
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-16 09:44:17 -03:00
commit 8f23c29818

View file

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