Refactor APIRequest and TextToRecord components

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-07 15:59:13 -03:00
commit a01a22ce3e
5 changed files with 39 additions and 65 deletions

View file

@ -1,12 +1,11 @@
import asyncio import asyncio
import json
from typing import List, Optional from typing import List, Optional
import httpx import httpx
import json
from langflow import CustomComponent from langflow import CustomComponent
from langflow.schema import Record from langflow.schema import Record
from langflow.services.database.models.base import orjson_dumps
class APIRequest(CustomComponent): class APIRequest(CustomComponent):
@ -52,35 +51,44 @@ class APIRequest(CustomComponent):
timeout: int = 5, timeout: int = 5,
) -> Record: ) -> Record:
method = method.upper() method = method.upper()
if method not in ["GET", "POST", "PATCH", "PUT"]: if method not in ["GET", "POST", "PATCH", "PUT", "DELETE"]:
raise ValueError(f"Unsupported method: {method}") raise ValueError(f"Unsupported method: {method}")
data = body if body else None data = body if body else None
data = json.dumps(data) data = json.dumps(data)
try: try:
response = await client.request(method, url, headers=headers, content=data, timeout=timeout) response = await client.request(
method, url, headers=headers, content=data, timeout=timeout
)
try: try:
response_json = response.json() result = response.json()
result = orjson_dumps(response_json, indent_2=False)
except Exception: except Exception:
result = response.text result = response.text
return Record( return Record(
text=result,
data={ data={
"source": url, "source": url,
"headers": headers, "headers": headers,
"status_code": response.status_code, "status_code": response.status_code,
"result": result,
}, },
) )
except httpx.TimeoutException: except httpx.TimeoutException:
return Record( return Record(
text="Request Timed Out", data={
data={"source": url, "headers": headers, "status_code": 408}, "source": url,
"headers": headers,
"status_code": 408,
"error": "Request timed out",
},
) )
except Exception as exc: except Exception as exc:
return Record( return Record(
text=str(exc), data={
data={"source": url, "headers": headers, "status_code": 500}, "source": url,
"headers": headers,
"status_code": 500,
"error": str(exc),
},
) )
async def build( async def build(
@ -88,15 +96,23 @@ class APIRequest(CustomComponent):
method: str, method: str,
url: List[str], url: List[str],
headers: Optional[dict] = None, headers: Optional[dict] = None,
body: Optional[dict] = None, body: Optional[List[Record]] = None,
timeout: int = 5, timeout: int = 5,
) -> List[Record]: ) -> List[Record]:
if headers is None: if headers is None:
headers = {} headers = {}
urls = url if isinstance(url, list) else [url] urls = url if isinstance(url, list) else [url]
bodies = body if isinstance(body, list) else [body] if body else [None] * len(urls) bodies = []
if body:
if isinstance(body, list):
bodies = [b.data for b in body]
else:
bodies = [body.data]
async with httpx.AsyncClient() as client: async with httpx.AsyncClient() as client:
results = await asyncio.gather( results = await asyncio.gather(
*[self.make_request(client, method, u, headers, rec, timeout) for u, rec in zip(urls, bodies)] *[
self.make_request(client, method, u, headers, rec, timeout)
for u, rec in zip(urls, bodies)
]
) )
return results return results

View file

@ -80,7 +80,7 @@ class TextToRecordComponent(CustomComponent):
"display_name": "Mode", "display_name": "Mode",
"options": ["Text", "Number"], "options": ["Text", "Number"],
"info": "The mode to use for creating the record.", "info": "The mode to use for creating the record.",
"refresh": True, "real_time_refresh": True,
"input_types": [], "input_types": [],
}, },
} }

View file

@ -9,7 +9,9 @@ class UUIDGeneratorComponent(CustomComponent):
display_name = "Unique ID Generator" display_name = "Unique ID Generator"
description = "Generates a unique ID." description = "Generates a unique ID."
def update_build_config(self, build_config: dict, field_name: Text, field_value: Any): def update_build_config(
self, build_config: dict, field_name: Text, field_value: Any
):
if field_name == "unique_id": if field_name == "unique_id":
build_config[field_name]["value"] = str(uuid.uuid4()) build_config[field_name]["value"] = str(uuid.uuid4())
return build_config return build_config
@ -18,7 +20,7 @@ class UUIDGeneratorComponent(CustomComponent):
return { return {
"unique_id": { "unique_id": {
"display_name": "Value", "display_name": "Value",
"refresh": True, "real_time_refresh": True,
} }
} }

View file

@ -1,47 +0,0 @@
from typing import Any
from langflow import CustomComponent
from langflow.schema import Record
from langflow.template.field.base import TemplateField
class RecordComponent(CustomComponent):
display_name = "Record Numbers"
description = "A component to create a record from key-value pairs."
field_order = ["n_keys"]
def update_build_config(self, build_config: dict, field_name: str, field_value: Any):
if field_value is None:
return
elif int(field_value) == 0:
keep = ["n_keys", "code"]
for key in build_config.copy():
if key in keep:
continue
del build_config[key]
build_config[field_name]["value"] = int(field_value)
# Add new fields depending on the field value
for i in range(int(field_value)):
field = TemplateField(
name=f"Key and Value {i}",
field_type="dict",
display_name="",
info="The key for the record.",
input_types=["Text"],
)
build_config[field.name] = field.to_dict()
def build_config(self):
return {
"n_keys": {
"display_name": "Number of Fields",
"refresh": True,
"info": "The number of keys to create in the record.",
},
}
def build(self, n_keys: int, **kwargs) -> Record:
data = {k: v for d in kwargs.values() for k, v in d.items()}
record = Record(data=data)
return record

View file

@ -140,7 +140,10 @@ class CustomComponent(Component):
return self.field_config return self.field_config
def update_build_config( def update_build_config(
self, build_config: dotdict, field_name: str, field_value: Any self,
build_config: dotdict,
field_name: str,
field_value: Any,
): ):
build_config[field_name] = field_value build_config[field_name] = field_value
return build_config return build_config