Refactor process.py for readability and maintainability
This commit is contained in:
parent
c6045809f2
commit
4b14aef8a7
1 changed files with 31 additions and 11 deletions
|
|
@ -126,7 +126,9 @@ async def process_runnable(runnable: Runnable, inputs: Union[dict, List[dict]]):
|
||||||
elif isinstance(inputs, dict) and hasattr(runnable, "ainvoke"):
|
elif isinstance(inputs, dict) and hasattr(runnable, "ainvoke"):
|
||||||
result = await runnable.ainvoke(inputs)
|
result = await runnable.ainvoke(inputs)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Runnable {runnable} does not support inputs of type {type(inputs)}")
|
raise ValueError(
|
||||||
|
f"Runnable {runnable} does not support inputs of type {type(inputs)}"
|
||||||
|
)
|
||||||
# Check if the result is a list of AIMessages
|
# Check if the result is a list of AIMessages
|
||||||
if isinstance(result, list) and all(isinstance(r, AIMessage) for r in result):
|
if isinstance(result, list) and all(isinstance(r, AIMessage) for r in result):
|
||||||
result = [r.content for r in result]
|
result = [r.content for r in result]
|
||||||
|
|
@ -135,7 +137,9 @@ async def process_runnable(runnable: Runnable, inputs: Union[dict, List[dict]]):
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
async def process_inputs_dict(built_object: Union[Chain, VectorStore, Runnable], inputs: dict):
|
async def process_inputs_dict(
|
||||||
|
built_object: Union[Chain, VectorStore, Runnable], inputs: dict
|
||||||
|
):
|
||||||
if isinstance(built_object, Chain):
|
if isinstance(built_object, Chain):
|
||||||
if inputs is None:
|
if inputs is None:
|
||||||
raise ValueError("Inputs must be provided for a Chain")
|
raise ValueError("Inputs must be provided for a Chain")
|
||||||
|
|
@ -170,7 +174,9 @@ async def process_inputs_list(built_object: Runnable, inputs: List[dict]):
|
||||||
return await process_runnable(built_object, inputs)
|
return await process_runnable(built_object, inputs)
|
||||||
|
|
||||||
|
|
||||||
async def generate_result(built_object: Union[Chain, VectorStore, Runnable], inputs: Union[dict, List[dict]]):
|
async def generate_result(
|
||||||
|
built_object: Union[Chain, VectorStore, Runnable], inputs: Union[dict, List[dict]]
|
||||||
|
):
|
||||||
if isinstance(inputs, dict):
|
if isinstance(inputs, dict):
|
||||||
result = await process_inputs_dict(built_object, inputs)
|
result = await process_inputs_dict(built_object, inputs)
|
||||||
elif isinstance(inputs, List) and isinstance(built_object, Runnable):
|
elif isinstance(inputs, List) and isinstance(built_object, Runnable):
|
||||||
|
|
@ -208,24 +214,30 @@ async def run_graph(
|
||||||
else:
|
else:
|
||||||
graph_data = graph._graph_data
|
graph_data = graph._graph_data
|
||||||
if not session_id and session_service is not None:
|
if not session_id and session_service is not None:
|
||||||
session_id = session_service.generate_key(session_id=flow_id, data_graph=graph_data)
|
session_id = session_service.generate_key(
|
||||||
|
session_id=flow_id, data_graph=graph_data
|
||||||
|
)
|
||||||
if inputs is None:
|
if inputs is None:
|
||||||
inputs = {}
|
inputs = {}
|
||||||
|
|
||||||
outputs = await graph.run(inputs, stream=stream)
|
outputs = await graph.run(inputs, stream=stream, session_id=session_id)
|
||||||
if session_id and session_service:
|
if session_id and session_service:
|
||||||
session_service.update_session(session_id, (graph, artifacts))
|
session_service.update_session(session_id, (graph, artifacts))
|
||||||
return outputs, session_id
|
return outputs, session_id
|
||||||
|
|
||||||
|
|
||||||
def validate_input(graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]) -> List[Dict[str, Any]]:
|
def validate_input(
|
||||||
|
graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]
|
||||||
|
) -> List[Dict[str, Any]]:
|
||||||
if not isinstance(graph_data, dict) or not isinstance(tweaks, dict):
|
if not isinstance(graph_data, dict) or not isinstance(tweaks, dict):
|
||||||
raise ValueError("graph_data and tweaks should be dictionaries")
|
raise ValueError("graph_data and tweaks should be dictionaries")
|
||||||
|
|
||||||
nodes = graph_data.get("data", {}).get("nodes") or graph_data.get("nodes")
|
nodes = graph_data.get("data", {}).get("nodes") or graph_data.get("nodes")
|
||||||
|
|
||||||
if not isinstance(nodes, list):
|
if not isinstance(nodes, list):
|
||||||
raise ValueError("graph_data should contain a list of nodes under 'data' key or directly under 'nodes' key")
|
raise ValueError(
|
||||||
|
"graph_data should contain a list of nodes under 'data' key or directly under 'nodes' key"
|
||||||
|
)
|
||||||
|
|
||||||
return nodes
|
return nodes
|
||||||
|
|
||||||
|
|
@ -234,7 +246,9 @@ def apply_tweaks(node: Dict[str, Any], node_tweaks: Dict[str, Any]) -> None:
|
||||||
template_data = node.get("data", {}).get("node", {}).get("template")
|
template_data = node.get("data", {}).get("node", {}).get("template")
|
||||||
|
|
||||||
if not isinstance(template_data, dict):
|
if not isinstance(template_data, dict):
|
||||||
logger.warning(f"Template data for node {node.get('id')} should be a dictionary")
|
logger.warning(
|
||||||
|
f"Template data for node {node.get('id')} should be a dictionary"
|
||||||
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
for tweak_name, tweak_value in node_tweaks.items():
|
for tweak_name, tweak_value in node_tweaks.items():
|
||||||
|
|
@ -249,7 +263,9 @@ def apply_tweaks_on_vertex(vertex: Vertex, node_tweaks: Dict[str, Any]) -> None:
|
||||||
vertex.params[tweak_name] = tweak_value
|
vertex.params[tweak_name] = tweak_value
|
||||||
|
|
||||||
|
|
||||||
def process_tweaks(graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
|
def process_tweaks(
|
||||||
|
graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]
|
||||||
|
) -> Dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
This function is used to tweak the graph data using the node id and the tweaks dict.
|
This function is used to tweak the graph data using the node id and the tweaks dict.
|
||||||
|
|
||||||
|
|
@ -270,7 +286,9 @@ def process_tweaks(graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]
|
||||||
if node_tweaks := tweaks.get(node_id):
|
if node_tweaks := tweaks.get(node_id):
|
||||||
apply_tweaks(node, node_tweaks)
|
apply_tweaks(node, node_tweaks)
|
||||||
else:
|
else:
|
||||||
logger.warning("Each node should be a dictionary with an 'id' key of type str")
|
logger.warning(
|
||||||
|
"Each node should be a dictionary with an 'id' key of type str"
|
||||||
|
)
|
||||||
|
|
||||||
return graph_data
|
return graph_data
|
||||||
|
|
||||||
|
|
@ -282,6 +300,8 @@ def process_tweaks_on_graph(graph: Graph, tweaks: Dict[str, Dict[str, Any]]):
|
||||||
if node_tweaks := tweaks.get(node_id):
|
if node_tweaks := tweaks.get(node_id):
|
||||||
apply_tweaks_on_vertex(vertex, node_tweaks)
|
apply_tweaks_on_vertex(vertex, node_tweaks)
|
||||||
else:
|
else:
|
||||||
logger.warning("Each node should be a Vertex with an 'id' attribute of type str")
|
logger.warning(
|
||||||
|
"Each node should be a Vertex with an 'id' attribute of type str"
|
||||||
|
)
|
||||||
|
|
||||||
return graph
|
return graph
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue