Fix styleUtils import and remove unnecessary lines
This commit is contained in:
parent
3a07ad38b0
commit
9cfb03fbc9
47 changed files with 219 additions and 574 deletions
|
|
@ -3,8 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
|||
from uuid import UUID
|
||||
|
||||
from langchain.schema import AgentAction, AgentFinish
|
||||
from langchain_core.callbacks.base import (AsyncCallbackHandler,
|
||||
BaseCallbackHandler)
|
||||
from langchain_core.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
|
||||
from langflow.api.v1.schemas import ChatResponse, PromptResponse
|
||||
from langflow.services.deps import get_chat_service
|
||||
from langflow.utils.util import remove_ansi_escape_codes
|
||||
|
|
|
|||
|
|
@ -43,13 +43,9 @@ async def chat(
|
|||
user = await get_current_user_for_websocket(websocket, db)
|
||||
await websocket.accept()
|
||||
if not user:
|
||||
await websocket.close(
|
||||
code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized"
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized")
|
||||
elif not user.is_active:
|
||||
await websocket.close(
|
||||
code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized"
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized")
|
||||
|
||||
if client_id in chat_service.cache_service:
|
||||
await chat_service.handle_websocket(client_id, websocket)
|
||||
|
|
@ -65,9 +61,7 @@ async def chat(
|
|||
logger.error(f"Error in chat websocket: {exc}")
|
||||
messsage = exc.detail if isinstance(exc, HTTPException) else str(exc)
|
||||
if "Could not validate credentials" in str(exc):
|
||||
await websocket.close(
|
||||
code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized"
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized")
|
||||
else:
|
||||
await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=messsage)
|
||||
|
||||
|
|
@ -137,12 +131,8 @@ async def build_vertex(
|
|||
cache = chat_service.get_cache(flow_id)
|
||||
if not cache:
|
||||
# If there's no cache
|
||||
logger.warning(
|
||||
f"No cache found for {flow_id}. Building graph starting at {vertex_id}"
|
||||
)
|
||||
graph = build_and_cache_graph(
|
||||
flow_id=flow_id, session=next(get_session()), chat_service=chat_service
|
||||
)
|
||||
logger.warning(f"No cache found for {flow_id}. Building graph starting at {vertex_id}")
|
||||
graph = build_and_cache_graph(flow_id=flow_id, session=next(get_session()), chat_service=chat_service)
|
||||
else:
|
||||
graph = cache.get("result")
|
||||
result_dict = {}
|
||||
|
|
|
|||
|
|
@ -66,8 +66,6 @@ async def get_transactions(
|
|||
monitor_service: MonitorService = Depends(get_monitor_service),
|
||||
):
|
||||
try:
|
||||
return monitor_service.get_transactions(
|
||||
source=source, target=target, status=status, order_by=order_by
|
||||
)
|
||||
return monitor_service.get_transactions(source=source, target=target, status=status, order_by=order_by)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
|
|
|||
|
|
@ -161,9 +161,7 @@ class StreamData(BaseModel):
|
|||
data: dict
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
f"event: {self.event}\ndata: {orjson_dumps(self.data, indent_2=False)}\n\n"
|
||||
)
|
||||
return f"event: {self.event}\ndata: {orjson_dumps(self.data, indent_2=False)}\n\n"
|
||||
|
||||
|
||||
class CustomComponentCode(BaseModel):
|
||||
|
|
|
|||
|
|
@ -40,9 +40,7 @@ def post_validate_prompt(prompt_request: ValidatePromptRequest):
|
|||
|
||||
add_new_variables_to_template(input_variables, prompt_request)
|
||||
|
||||
remove_old_variables_from_template(
|
||||
old_custom_fields, input_variables, prompt_request
|
||||
)
|
||||
remove_old_variables_from_template(old_custom_fields, input_variables, prompt_request)
|
||||
|
||||
update_input_variables_field(input_variables, prompt_request)
|
||||
|
||||
|
|
@ -57,19 +55,12 @@ def post_validate_prompt(prompt_request: ValidatePromptRequest):
|
|||
|
||||
def get_old_custom_fields(prompt_request):
|
||||
try:
|
||||
if (
|
||||
len(prompt_request.frontend_node.custom_fields) == 1
|
||||
and prompt_request.name == ""
|
||||
):
|
||||
if len(prompt_request.frontend_node.custom_fields) == 1 and prompt_request.name == "":
|
||||
# If there is only one custom field and the name is empty string
|
||||
# then we are dealing with the first prompt request after the node was created
|
||||
prompt_request.name = list(
|
||||
prompt_request.frontend_node.custom_fields.keys()
|
||||
)[0]
|
||||
prompt_request.name = list(prompt_request.frontend_node.custom_fields.keys())[0]
|
||||
|
||||
old_custom_fields = prompt_request.frontend_node.custom_fields[
|
||||
prompt_request.name
|
||||
]
|
||||
old_custom_fields = prompt_request.frontend_node.custom_fields[prompt_request.name]
|
||||
if old_custom_fields is None:
|
||||
old_custom_fields = []
|
||||
|
||||
|
|
@ -95,40 +86,26 @@ def add_new_variables_to_template(input_variables, prompt_request):
|
|||
)
|
||||
if variable in prompt_request.frontend_node.template:
|
||||
# Set the new field with the old value
|
||||
template_field.value = prompt_request.frontend_node.template[variable][
|
||||
"value"
|
||||
]
|
||||
template_field.value = prompt_request.frontend_node.template[variable]["value"]
|
||||
|
||||
prompt_request.frontend_node.template[variable] = template_field.to_dict()
|
||||
|
||||
# Check if variable is not already in the list before appending
|
||||
if (
|
||||
variable
|
||||
not in prompt_request.frontend_node.custom_fields[prompt_request.name]
|
||||
):
|
||||
prompt_request.frontend_node.custom_fields[prompt_request.name].append(
|
||||
variable
|
||||
)
|
||||
if variable not in prompt_request.frontend_node.custom_fields[prompt_request.name]:
|
||||
prompt_request.frontend_node.custom_fields[prompt_request.name].append(variable)
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception(exc)
|
||||
raise HTTPException(status_code=500, detail=str(exc)) from exc
|
||||
|
||||
|
||||
def remove_old_variables_from_template(
|
||||
old_custom_fields, input_variables, prompt_request
|
||||
):
|
||||
def remove_old_variables_from_template(old_custom_fields, input_variables, prompt_request):
|
||||
for variable in old_custom_fields:
|
||||
if variable not in input_variables:
|
||||
try:
|
||||
# Remove the variable from custom_fields associated with the given name
|
||||
if (
|
||||
variable
|
||||
in prompt_request.frontend_node.custom_fields[prompt_request.name]
|
||||
):
|
||||
prompt_request.frontend_node.custom_fields[
|
||||
prompt_request.name
|
||||
].remove(variable)
|
||||
if variable in prompt_request.frontend_node.custom_fields[prompt_request.name]:
|
||||
prompt_request.frontend_node.custom_fields[prompt_request.name].remove(variable)
|
||||
|
||||
# Remove the variable from the template
|
||||
prompt_request.frontend_node.template.pop(variable, None)
|
||||
|
|
@ -140,6 +117,4 @@ def remove_old_variables_from_template(
|
|||
|
||||
def update_input_variables_field(input_variables, prompt_request):
|
||||
if "input_variables" in prompt_request.frontend_node.template:
|
||||
prompt_request.frontend_node.template["input_variables"][
|
||||
"value"
|
||||
] = input_variables
|
||||
prompt_request.frontend_node.template["input_variables"]["value"] = input_variables
|
||||
|
|
|
|||
|
|
@ -52,9 +52,7 @@ class StoreMessages(CustomComponent):
|
|||
if not records:
|
||||
records = []
|
||||
if not session_id or not sender or not sender_name:
|
||||
raise ValueError(
|
||||
"If passing texts, session_id, sender, and sender_name must be provided."
|
||||
)
|
||||
raise ValueError("If passing texts, session_id, sender, and sender_name must be provided.")
|
||||
for text in texts:
|
||||
record = Record(
|
||||
text=text,
|
||||
|
|
|
|||
|
|
@ -126,8 +126,7 @@ class ChatLiteLLMComponent(CustomComponent):
|
|||
litellm.set_verbose = verbose
|
||||
except ImportError:
|
||||
raise ChatLiteLLMException(
|
||||
"Could not import litellm python package. "
|
||||
"Please install it with `pip install litellm`"
|
||||
"Could not import litellm python package. " "Please install it with `pip install litellm`"
|
||||
)
|
||||
provider_map = {
|
||||
"OpenAI": "openai_api_key",
|
||||
|
|
|
|||
|
|
@ -63,4 +63,4 @@ class AmazonBedrockComponent(CustomComponent):
|
|||
message = output.invoke(inputs)
|
||||
result = message.content if hasattr(message, "content") else message
|
||||
self.status = result
|
||||
return result
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ class AnthropicLLM(CustomComponent):
|
|||
def build(
|
||||
self,
|
||||
model: str,
|
||||
inputs:str,
|
||||
inputs: str,
|
||||
anthropic_api_key: Optional[str] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
|
|
@ -78,4 +78,3 @@ class AnthropicLLM(CustomComponent):
|
|||
result = message.content if hasattr(message, "content") else message
|
||||
self.status = result
|
||||
return result
|
||||
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ class CTransformersComponent(CustomComponent):
|
|||
"inputs": {"display_name": "Input"},
|
||||
}
|
||||
|
||||
def build(self, model: str, model_file: str,inputs:str, model_type: str, config: Optional[Dict] = None) -> Text:
|
||||
def build(self, model: str, model_file: str, inputs: str, model_type: str, config: Optional[Dict] = None) -> Text:
|
||||
output = CTransformers(model=model, model_file=model_file, model_type=model_type, config=config)
|
||||
message = output.invoke(inputs)
|
||||
result = message.content if hasattr(message, "content") else message
|
||||
|
|
|
|||
|
|
@ -14,41 +14,41 @@ class GoogleGenerativeAIComponent(CustomComponent):
|
|||
|
||||
def build_config(self):
|
||||
return {
|
||||
"google_api_key":
|
||||
{ "display_name":"Google API Key",
|
||||
"info":"The Google API Key to use for the Google Generative AI.",
|
||||
} ,
|
||||
"max_output_tokens":{
|
||||
"display_name":"Max Output Tokens",
|
||||
"info":"The maximum number of tokens to generate.",
|
||||
"google_api_key": {
|
||||
"display_name": "Google API Key",
|
||||
"info": "The Google API Key to use for the Google Generative AI.",
|
||||
},
|
||||
"max_output_tokens": {
|
||||
"display_name": "Max Output Tokens",
|
||||
"info": "The maximum number of tokens to generate.",
|
||||
},
|
||||
"temperature": {
|
||||
"display_name":"Temperature",
|
||||
"info":"Run inference with this temperature. Must by in the closed interval [0.0, 1.0].",
|
||||
"display_name": "Temperature",
|
||||
"info": "Run inference with this temperature. Must by in the closed interval [0.0, 1.0].",
|
||||
},
|
||||
"top_k": {
|
||||
"display_name":"Top K",
|
||||
"info":"Decode using top-k sampling: consider the set of top_k most probable tokens. Must be positive.",
|
||||
"range_spec":RangeSpec(min=0, max=2, step=0.1),
|
||||
"advanced":True,
|
||||
"display_name": "Top K",
|
||||
"info": "Decode using top-k sampling: consider the set of top_k most probable tokens. Must be positive.",
|
||||
"range_spec": RangeSpec(min=0, max=2, step=0.1),
|
||||
"advanced": True,
|
||||
},
|
||||
"top_p": {
|
||||
"display_name":"Top P",
|
||||
"info":"The maximum cumulative probability of tokens to consider when sampling.",
|
||||
"advanced":True,
|
||||
"display_name": "Top P",
|
||||
"info": "The maximum cumulative probability of tokens to consider when sampling.",
|
||||
"advanced": True,
|
||||
},
|
||||
"n": {
|
||||
"display_name":"N",
|
||||
"info":"Number of chat completions to generate for each prompt. Note that the API may not return the full n completions if duplicates are generated.",
|
||||
"advanced":True,
|
||||
"display_name": "N",
|
||||
"info": "Number of chat completions to generate for each prompt. Note that the API may not return the full n completions if duplicates are generated.",
|
||||
"advanced": True,
|
||||
},
|
||||
"model": {
|
||||
"display_name":"Model",
|
||||
"info":"The name of the model to use. Supported examples: gemini-pro",
|
||||
"options":["gemini-pro", "gemini-pro-vision"],
|
||||
"display_name": "Model",
|
||||
"info": "The name of the model to use. Supported examples: gemini-pro",
|
||||
"options": ["gemini-pro", "gemini-pro-vision"],
|
||||
},
|
||||
"code": {
|
||||
"advanced":True,
|
||||
"advanced": True,
|
||||
},
|
||||
"inputs": {"display_name": "Input"},
|
||||
}
|
||||
|
|
@ -57,14 +57,14 @@ class GoogleGenerativeAIComponent(CustomComponent):
|
|||
self,
|
||||
google_api_key: str,
|
||||
model: str,
|
||||
inputs:str,
|
||||
inputs: str,
|
||||
max_output_tokens: Optional[int] = None,
|
||||
temperature: float = 0.1,
|
||||
top_k: Optional[int] = None,
|
||||
top_p: Optional[float] = None,
|
||||
n: Optional[int] = 1,
|
||||
) -> Text:
|
||||
output = ChatGoogleGenerativeAI(
|
||||
output = ChatGoogleGenerativeAI(
|
||||
model=model,
|
||||
max_output_tokens=max_output_tokens or None, # type: ignore
|
||||
temperature=temperature,
|
||||
|
|
|
|||
|
|
@ -47,4 +47,3 @@ class HuggingFaceEndpointsComponent(CustomComponent):
|
|||
result = message.content if hasattr(message, "content") else message
|
||||
self.status = result
|
||||
return result
|
||||
|
||||
|
|
|
|||
|
|
@ -57,7 +57,7 @@ class LlamaCppComponent(CustomComponent):
|
|||
def build(
|
||||
self,
|
||||
model_path: str,
|
||||
inputs:str,
|
||||
inputs: str,
|
||||
grammar: Optional[str] = None,
|
||||
cache: Optional[bool] = None,
|
||||
client: Optional[Any] = None,
|
||||
|
|
|
|||
|
|
@ -171,7 +171,7 @@ class ChatOllamaComponent(CustomComponent):
|
|||
self,
|
||||
base_url: Optional[str],
|
||||
model: str,
|
||||
inputs:str,
|
||||
inputs: str,
|
||||
mirostat: Optional[str],
|
||||
mirostat_eta: Optional[float] = None,
|
||||
mirostat_tau: Optional[float] = None,
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ class PromptComponent(CustomComponent):
|
|||
template: Prompt,
|
||||
**kwargs,
|
||||
) -> Text:
|
||||
|
||||
prompt_template = PromptTemplate.from_template(template)
|
||||
|
||||
attributes_to_check = ["text", "page_content"]
|
||||
|
|
|
|||
|
|
@ -27,9 +27,7 @@ class RecordsAsTextComponent(CustomComponent):
|
|||
if isinstance(records, Record):
|
||||
records = [records]
|
||||
|
||||
formated_records = [
|
||||
template.format(text=record.text, **record.data) for record in records
|
||||
]
|
||||
formated_records = [template.format(text=record.text, **record.data) for record in records]
|
||||
result_string = "\n".join(formated_records)
|
||||
self.status = result_string
|
||||
return result_string
|
||||
|
|
|
|||
|
|
@ -84,8 +84,7 @@ class ChromaComponent(CustomComponent):
|
|||
|
||||
if chroma_server_host is not None:
|
||||
chroma_settings = chromadb.config.Settings(
|
||||
chroma_server_cors_allow_origins=chroma_server_cors_allow_origins
|
||||
or None,
|
||||
chroma_server_cors_allow_origins=chroma_server_cors_allow_origins or None,
|
||||
chroma_server_host=chroma_server_host,
|
||||
chroma_server_port=chroma_server_port or None,
|
||||
chroma_server_grpc_port=chroma_server_grpc_port or None,
|
||||
|
|
@ -100,9 +99,7 @@ class ChromaComponent(CustomComponent):
|
|||
|
||||
if documents is not None and embedding is not None:
|
||||
if len(documents) == 0:
|
||||
raise ValueError(
|
||||
"If documents are provided, there must be at least one document."
|
||||
)
|
||||
raise ValueError("If documents are provided, there must be at least one document.")
|
||||
chroma = Chroma.from_documents(
|
||||
documents=documents, # type: ignore
|
||||
persist_directory=index_directory,
|
||||
|
|
@ -111,7 +108,5 @@ class ChromaComponent(CustomComponent):
|
|||
client_settings=chroma_settings,
|
||||
)
|
||||
else:
|
||||
chroma = Chroma(
|
||||
persist_directory=index_directory, client_settings=chroma_settings
|
||||
)
|
||||
chroma = Chroma(persist_directory=index_directory, client_settings=chroma_settings)
|
||||
return chroma
|
||||
|
|
|
|||
|
|
@ -92,8 +92,7 @@ class ChromaSearchComponent(CustomComponent):
|
|||
|
||||
if chroma_server_host is not None:
|
||||
chroma_settings = chromadb.config.Settings(
|
||||
chroma_server_cors_allow_origins=chroma_server_cors_allow_origins
|
||||
or None,
|
||||
chroma_server_cors_allow_origins=chroma_server_cors_allow_origins or None,
|
||||
chroma_server_host=chroma_server_host,
|
||||
chroma_server_port=chroma_server_port or None,
|
||||
chroma_server_grpc_port=chroma_server_grpc_port or None,
|
||||
|
|
|
|||
|
|
@ -11,9 +11,7 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
class SourceHandle(BaseModel):
|
||||
baseClasses: List[str] = Field(
|
||||
..., description="List of base classes for the source handle."
|
||||
)
|
||||
baseClasses: List[str] = Field(..., description="List of base classes for the source handle.")
|
||||
dataType: str = Field(..., description="Data type for the source handle.")
|
||||
id: str = Field(..., description="Unique identifier for the source handle.")
|
||||
|
||||
|
|
@ -21,9 +19,7 @@ class SourceHandle(BaseModel):
|
|||
class TargetHandle(BaseModel):
|
||||
fieldName: str = Field(..., description="Field name for the target handle.")
|
||||
id: str = Field(..., description="Unique identifier for the target handle.")
|
||||
inputTypes: Optional[List[str]] = Field(
|
||||
None, description="List of input types for the target handle."
|
||||
)
|
||||
inputTypes: Optional[List[str]] = Field(None, description="List of input types for the target handle.")
|
||||
type: str = Field(..., description="Type of the target handle.")
|
||||
|
||||
|
||||
|
|
@ -52,24 +48,16 @@ class Edge:
|
|||
|
||||
def validate_handles(self, source, target) -> None:
|
||||
if self.target_handle.inputTypes is None:
|
||||
self.valid_handles = (
|
||||
self.target_handle.type in self.source_handle.baseClasses
|
||||
)
|
||||
self.valid_handles = self.target_handle.type in self.source_handle.baseClasses
|
||||
else:
|
||||
self.valid_handles = (
|
||||
any(
|
||||
baseClass in self.target_handle.inputTypes
|
||||
for baseClass in self.source_handle.baseClasses
|
||||
)
|
||||
any(baseClass in self.target_handle.inputTypes for baseClass in self.source_handle.baseClasses)
|
||||
or self.target_handle.type in self.source_handle.baseClasses
|
||||
)
|
||||
if not self.valid_handles:
|
||||
logger.debug(self.source_handle)
|
||||
logger.debug(self.target_handle)
|
||||
raise ValueError(
|
||||
f"Edge between {source.vertex_type} and {target.vertex_type} "
|
||||
f"has invalid handles"
|
||||
)
|
||||
raise ValueError(f"Edge between {source.vertex_type} and {target.vertex_type} " f"has invalid handles")
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.source_id = state["source_id"]
|
||||
|
|
@ -86,11 +74,7 @@ class Edge:
|
|||
# Both lists contain strings and sometimes a string contains the value we are
|
||||
# looking for e.g. comgin_out=["Chain"] and target_reqs=["LLMChain"]
|
||||
# so we need to check if any of the strings in source_types is in target_reqs
|
||||
self.valid = any(
|
||||
output in target_req
|
||||
for output in self.source_types
|
||||
for target_req in self.target_reqs
|
||||
)
|
||||
self.valid = any(output in target_req for output in self.source_types for target_req in self.target_reqs)
|
||||
# Get what type of input the target node is expecting
|
||||
|
||||
self.matched_type = next(
|
||||
|
|
@ -101,10 +85,7 @@ class Edge:
|
|||
if no_matched_type:
|
||||
logger.debug(self.source_types)
|
||||
logger.debug(self.target_reqs)
|
||||
raise ValueError(
|
||||
f"Edge between {source.vertex_type} and {target.vertex_type} "
|
||||
f"has no matched type"
|
||||
)
|
||||
raise ValueError(f"Edge between {source.vertex_type} and {target.vertex_type} " f"has no matched type")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
|
|
@ -116,11 +97,7 @@ class Edge:
|
|||
return hash(self.__repr__())
|
||||
|
||||
def __eq__(self, __value: object) -> bool:
|
||||
return (
|
||||
self.__repr__() == __value.__repr__()
|
||||
if isinstance(__value, Edge)
|
||||
else False
|
||||
)
|
||||
return self.__repr__() == __value.__repr__() if isinstance(__value, Edge) else False
|
||||
|
||||
|
||||
class ContractEdge(Edge):
|
||||
|
|
@ -176,9 +153,7 @@ class ContractEdge(Edge):
|
|||
return f"{self.source_id} -[{self.target_param}]-> {self.target_id}"
|
||||
|
||||
|
||||
def log_transaction(
|
||||
edge: ContractEdge, source: "Vertex", target: "Vertex", status, error=None
|
||||
):
|
||||
def log_transaction(edge: ContractEdge, source: "Vertex", target: "Vertex", status, error=None):
|
||||
try:
|
||||
monitor_service = get_monitor_service()
|
||||
clean_params = build_clean_params(target)
|
||||
|
|
|
|||
|
|
@ -225,11 +225,7 @@ class Graph:
|
|||
return
|
||||
self.vertices.remove(vertex)
|
||||
self.vertex_map.pop(vertex_id)
|
||||
self.edges = [
|
||||
edge
|
||||
for edge in self.edges
|
||||
if edge.source_id != vertex_id and edge.target_id != vertex_id
|
||||
]
|
||||
self.edges = [edge for edge in self.edges if edge.source_id != vertex_id and edge.target_id != vertex_id]
|
||||
|
||||
def _build_vertex_params(self) -> None:
|
||||
"""Identifies and handles the LLM vertex within the graph."""
|
||||
|
|
@ -250,9 +246,7 @@ class Graph:
|
|||
return
|
||||
for vertex in self.vertices:
|
||||
if not self._validate_vertex(vertex):
|
||||
raise ValueError(
|
||||
f"{vertex.vertex_type} is not connected to any other components"
|
||||
)
|
||||
raise ValueError(f"{vertex.vertex_type} is not connected to any other components")
|
||||
|
||||
def _validate_vertex(self, vertex: Vertex) -> bool:
|
||||
"""Validates a vertex."""
|
||||
|
|
@ -268,11 +262,7 @@ class Graph:
|
|||
|
||||
def get_vertex_edges(self, vertex_id: str) -> List[ContractEdge]:
|
||||
"""Returns a list of edges for a given vertex."""
|
||||
return [
|
||||
edge
|
||||
for edge in self.edges
|
||||
if edge.source_id == vertex_id or edge.target_id == vertex_id
|
||||
]
|
||||
return [edge for edge in self.edges if edge.source_id == vertex_id or edge.target_id == vertex_id]
|
||||
|
||||
def get_vertices_with_target(self, vertex_id: str) -> List[Vertex]:
|
||||
"""Returns the vertices connected to a vertex."""
|
||||
|
|
@ -310,9 +300,7 @@ class Graph:
|
|||
def dfs(vertex):
|
||||
if state[vertex] == 1:
|
||||
# We have a cycle
|
||||
raise ValueError(
|
||||
"Graph contains a cycle, cannot perform topological sort"
|
||||
)
|
||||
raise ValueError("Graph contains a cycle, cannot perform topological sort")
|
||||
if state[vertex] == 0:
|
||||
state[vertex] = 1
|
||||
for edge in vertex.edges:
|
||||
|
|
@ -336,17 +324,11 @@ class Graph:
|
|||
|
||||
def get_predecessors(self, vertex):
|
||||
"""Returns the predecessors of a vertex."""
|
||||
return [
|
||||
self.get_vertex(source_id)
|
||||
for source_id in self.predecessor_map.get(vertex.id, [])
|
||||
]
|
||||
return [self.get_vertex(source_id) for source_id in self.predecessor_map.get(vertex.id, [])]
|
||||
|
||||
def get_successors(self, vertex):
|
||||
"""Returns the successors of a vertex."""
|
||||
return [
|
||||
self.get_vertex(target_id)
|
||||
for target_id in self.successor_map.get(vertex.id, [])
|
||||
]
|
||||
return [self.get_vertex(target_id) for target_id in self.successor_map.get(vertex.id, [])]
|
||||
|
||||
def get_vertex_neighbors(self, vertex: Vertex) -> Dict[Vertex, int]:
|
||||
"""Returns the neighbors of a vertex."""
|
||||
|
|
@ -385,9 +367,7 @@ class Graph:
|
|||
edges.append(ContractEdge(source, target, edge))
|
||||
return edges
|
||||
|
||||
def _get_vertex_class(
|
||||
self, node_type: str, node_base_type: str, node_id: str
|
||||
) -> Type[Vertex]:
|
||||
def _get_vertex_class(self, node_type: str, node_base_type: str, node_id: str) -> Type[Vertex]:
|
||||
"""Returns the node class based on the node type."""
|
||||
# First we check for the node_base_type
|
||||
node_name = node_id.split("-")[0]
|
||||
|
|
@ -417,18 +397,14 @@ class Graph:
|
|||
vertex_type: str = vertex_data["type"] # type: ignore
|
||||
vertex_base_type: str = vertex_data["node"]["template"]["_type"] # type: ignore
|
||||
|
||||
VertexClass = self._get_vertex_class(
|
||||
vertex_type, vertex_base_type, vertex_data["id"]
|
||||
)
|
||||
VertexClass = self._get_vertex_class(vertex_type, vertex_base_type, vertex_data["id"])
|
||||
vertex_instance = VertexClass(vertex, graph=self)
|
||||
vertex_instance.set_top_level(self.top_level_vertices)
|
||||
vertices.append(vertex_instance)
|
||||
|
||||
return vertices
|
||||
|
||||
def get_children_by_vertex_type(
|
||||
self, vertex: Vertex, vertex_type: str
|
||||
) -> List[Vertex]:
|
||||
def get_children_by_vertex_type(self, vertex: Vertex, vertex_type: str) -> List[Vertex]:
|
||||
"""Returns the children of a vertex based on the vertex type."""
|
||||
children = []
|
||||
vertex_types = [vertex.data["type"]]
|
||||
|
|
@ -440,9 +416,7 @@ class Graph:
|
|||
|
||||
def __repr__(self):
|
||||
vertex_ids = [vertex.id for vertex in self.vertices]
|
||||
edges_repr = "\n".join(
|
||||
[f"{edge.source_id} --> {edge.target_id}" for edge in self.edges]
|
||||
)
|
||||
edges_repr = "\n".join([f"{edge.source_id} --> {edge.target_id}" for edge in self.edges])
|
||||
return f"Graph:\nNodes: {vertex_ids}\nConnections:\n{edges_repr}"
|
||||
|
||||
def sort_up_to_vertex(self, vertex_id: str) -> "Graph":
|
||||
|
|
@ -473,9 +447,7 @@ class Graph:
|
|||
"""Performs a layered topological sort of the vertices in the graph."""
|
||||
|
||||
# Queue for vertices with no incoming edges
|
||||
queue = deque(
|
||||
vertex.id for vertex in vertices if self.in_degree_map[vertex.id] == 0
|
||||
)
|
||||
queue = deque(vertex.id for vertex in vertices if self.in_degree_map[vertex.id] == 0)
|
||||
layers = []
|
||||
|
||||
current_layer = 0
|
||||
|
|
@ -531,10 +503,7 @@ class Graph:
|
|||
|
||||
return refined_layers
|
||||
|
||||
def sort_chat_inputs_first(
|
||||
self, vertices_layers: List[List[str]]
|
||||
) -> List[List[str]]:
|
||||
|
||||
def sort_chat_inputs_first(self, vertices_layers: List[List[str]]) -> List[List[str]]:
|
||||
chat_inputs_first = []
|
||||
for layer in vertices_layers:
|
||||
for vertex_id in layer:
|
||||
|
|
@ -561,15 +530,11 @@ class Graph:
|
|||
self.increment_run_count()
|
||||
return vertices_layers
|
||||
|
||||
def sort_interface_components_first(
|
||||
self, vertices_layers: List[List[str]]
|
||||
) -> List[List[str]]:
|
||||
def sort_interface_components_first(self, vertices_layers: List[List[str]]) -> List[List[str]]:
|
||||
"""Sorts the vertices in the graph so that vertices containing ChatInput or ChatOutput come first."""
|
||||
|
||||
def contains_interface_component(vertex):
|
||||
return any(
|
||||
component.value in vertex for component in InterfaceComponentTypes
|
||||
)
|
||||
return any(component.value in vertex for component in InterfaceComponentTypes)
|
||||
|
||||
# Sort each inner list so that vertices containing ChatInput or ChatOutput come first
|
||||
sorted_vertices = [
|
||||
|
|
@ -588,13 +553,9 @@ class Graph:
|
|||
"""Sorts the vertices in the graph so that vertices with the lowest average build time come first."""
|
||||
if len(vertices_ids) == 1:
|
||||
return vertices_ids
|
||||
vertices_ids.sort(
|
||||
key=lambda vertex_id: self.get_vertex(vertex_id).avg_build_time
|
||||
)
|
||||
vertices_ids.sort(key=lambda vertex_id: self.get_vertex(vertex_id).avg_build_time)
|
||||
|
||||
return vertices_ids
|
||||
|
||||
sorted_vertices = [
|
||||
sort_layer_by_avg_build_time(layer) for layer in vertices_layers
|
||||
]
|
||||
sorted_vertices = [sort_layer_by_avg_build_time(layer) for layer in vertices_layers]
|
||||
return sorted_vertices
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue