Refactor process.py and schema.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-23 00:21:27 -03:00
commit dd9347a186
2 changed files with 7 additions and 76 deletions

View file

@ -1,20 +1,15 @@
import asyncio from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Coroutine, Dict, List, Optional, Tuple, Union
from langchain.agents import AgentExecutor from langchain.agents import AgentExecutor
from langchain.chains.base import Chain from langchain.schema import AgentAction
from langchain.schema import AgentAction, Document
from langchain_community.vectorstores import VectorStore
from langchain_core.messages import AIMessage
from langchain_core.runnables.base import Runnable
from loguru import logger from loguru import logger
from pydantic import BaseModel from pydantic import BaseModel
from langflow.graph.graph.base import Graph from langflow.graph.graph.base import Graph
from langflow.graph.schema import INPUT_FIELD_NAME, RunOutputs from langflow.graph.schema import RunOutputs
from langflow.graph.vertex.base import Vertex from langflow.graph.vertex.base import Vertex
from langflow.interface.custom.custom_component import CustomComponent
from langflow.interface.run import get_memory_key, update_memory_keys from langflow.interface.run import get_memory_key, update_memory_keys
from langflow.schema.schema import INPUT_FIELD_NAME
from langflow.services.session.service import SessionService from langflow.services.session.service import SessionService
if TYPE_CHECKING: if TYPE_CHECKING:
@ -124,73 +119,6 @@ def update_inputs_dict(inputs: dict, artifacts: Dict[str, Any]) -> dict:
return inputs return inputs
async def process_runnable(runnable: Runnable, inputs: Union[dict, List[dict]]):
if isinstance(inputs, List) and hasattr(runnable, "abatch"):
result = await runnable.abatch(inputs)
elif isinstance(inputs, dict) and hasattr(runnable, "ainvoke"):
result = await runnable.ainvoke(inputs)
else:
raise ValueError(f"Runnable {runnable} does not support inputs of type {type(inputs)}")
# Check if the result is a list of AIMessages
if isinstance(result, list) and all(isinstance(r, AIMessage) for r in result):
result = [r.content for r in result]
elif isinstance(result, AIMessage):
result = result.content
return result
async def process_inputs_dict(built_object: Union[Chain, VectorStore, Runnable], inputs: dict):
if isinstance(built_object, Chain):
if inputs is None:
raise ValueError("Inputs must be provided for a Chain")
logger.debug("Generating result and thought")
result = get_result_and_thought(built_object, inputs)
logger.debug("Generated result and thought")
elif isinstance(built_object, VectorStore) and "query" in inputs:
if isinstance(inputs, dict) and "search_type" not in inputs:
inputs["search_type"] = "similarity"
logger.info("search_type not provided, using default value: similarity")
result = built_object.search(**inputs)
elif isinstance(built_object, Document):
result = built_object.dict()
elif isinstance(built_object, Runnable):
result = await process_runnable(built_object, inputs)
if isinstance(result, list):
result = [r.content if hasattr(r, "content") else r for r in result]
elif hasattr(result, "content"):
result = result.content
else:
result = result
elif hasattr(built_object, "run") and isinstance(built_object, CustomComponent):
result = built_object.run(inputs)
else:
result = None
return result
async def process_inputs_list(built_object: Runnable, inputs: List[dict]):
return await process_runnable(built_object, inputs)
async def generate_result(built_object: Union[Chain, VectorStore, Runnable], inputs: Union[dict, List[dict]]):
if isinstance(inputs, dict):
result = await process_inputs_dict(built_object, inputs)
elif isinstance(inputs, List) and isinstance(built_object, Runnable):
result = await process_inputs_list(built_object, inputs)
else:
raise ValueError(f"Invalid inputs type: {type(inputs)}")
if result is None:
logger.warning(f"Unknown built_object type: {type(built_object)}")
if isinstance(built_object, Coroutine):
result = asyncio.run(built_object)
result = built_object
return result
class Result(BaseModel): class Result(BaseModel):
result: Any result: Any
session_id: str session_id: str

View file

@ -120,3 +120,6 @@ class Record(BaseModel):
# check which attributes the Record has by checking the keys in the data dictionary # check which attributes the Record has by checking the keys in the data dictionary
def __dir__(self): def __dir__(self):
return super().__dir__() + list(self.data.keys()) return super().__dir__() + list(self.data.keys())
INPUT_FIELD_NAME = "input_value"