Merge remote-tracking branch 'origin/dev' into celery

src/backend/langflow/api/v1/endpoints.py

@@ -1,12 +1,11 @@
 from http import HTTPStatus
-from typing import Annotated, Optional
+from typing import Annotated, Optional, Union
 
 from langflow.services.cache.utils import save_uploaded_file
 from langflow.services.database.models.flow import Flow
-from langflow.processing.process import process_tweaks
+from langflow.processing.process import process_graph_cached, process_tweaks
 from langflow.services.utils import get_settings_manager
 from langflow.utils.logger import logger
-from langflow.worker import process_graph_cached as process_graph_cached_worker
 from fastapi import APIRouter, Depends, HTTPException, UploadFile, Body
 
 from langflow.interface.custom.custom_component import CustomComponent
@@ -76,6 +75,7 @@ async def process_flow(
     inputs: Optional[dict] = None,
     tweaks: Optional[dict] = None,
     clear_cache: Annotated[bool, Body(embed=True)] = False,  # noqa: F821
+    session_id: Annotated[Union[None, str], Body(embed=True)] = None,  # noqa: F821
     session: Session = Depends(get_session),
 ):
     """
@@ -95,15 +95,10 @@ async def process_flow(
                 graph_data = process_tweaks(graph_data, tweaks)
             except Exception as exc:
                 logger.error(f"Error processing tweaks: {exc}")
-        # ! This was added just for testing purposes
-        response = process_graph_cached_worker.delay(
-            graph_data=graph_data,
-            inputs=inputs,
-            clear_cache=clear_cache,
-        ).get()
-        return ProcessResponse(
-            result=response,
+        response, session_id = process_graph_cached(
+            graph_data, inputs, clear_cache, session_id
         )
+        return ProcessResponse(result=response, session_id=session_id)
     except Exception as e:
         # Log stack trace
         logger.exception(e)

src/backend/langflow/api/v1/schemas.py

@@ -47,6 +47,7 @@ class ProcessResponse(BaseModel):
     """Process response schema."""
 
     result: dict
+    session_id: Optional[str] = None
 
 
 class ChatMessage(BaseModel):

src/backend/langflow/interface/run.py

@@ -1,3 +1,4 @@
+from typing import Any, Dict, Tuple
 from langflow.services.cache.utils import memoize_dict
 from langflow.graph import Graph
 from langflow.utils.logger import logger
@@ -15,7 +16,7 @@ def build_langchain_object_with_caching(data_graph):
 
 
 @memoize_dict(maxsize=10)
-def build_sorted_vertices_with_caching(data_graph):
+def build_sorted_vertices_with_caching(data_graph) -> Tuple[Any, Dict]:
     """
     Build langchain object from data_graph.
     """

src/backend/langflow/processing/process.py

@@ -85,39 +85,55 @@ def get_input_str_if_only_one_input(inputs: dict) -> Optional[str]:
     return list(inputs.values())[0] if len(inputs) == 1 else None
 
 
-def process_graph_cached(
-    data_graph: Dict[str, Any], inputs: Optional[dict] = None, clear_cache=False
-):
-    """
-    Process graph by extracting input variables and replacing ZeroShotPrompt
-    with PromptTemplate,then run the graph and return the result and thought.
-    """
-    # Load langchain object
+def get_build_result(data_graph, session_id):
+    # If session_id is provided, load the langchain_object from the session
+    # using build_sorted_vertices_with_caching.get_result_by_session_id
+    # if it returns something different than None, return it
+    # otherwise, build the graph and return the result
+    if session_id:
+        logger.debug(f"Loading LangChain object from session {session_id}")
+        result = build_sorted_vertices_with_caching.get_result_by_session_id(session_id)
+        if result is not None:
+            logger.debug("Loaded LangChain object")
+            return result
+
+    logger.debug("Building langchain object")
+    return build_sorted_vertices_with_caching(data_graph)
+
+
+def clear_caches_if_needed(clear_cache: bool):
     if clear_cache:
         build_sorted_vertices_with_caching.clear_cache()
         logger.debug("Cleared cache")
-    langchain_object, artifacts = build_sorted_vertices_with_caching(data_graph)
-    logger.debug("Loaded LangChain object")
-    if inputs is None:
-        inputs = {}
 
-    # Add artifacts to inputs
-    # artifacts can be documents loaded when building
-    # the flow
-    for (
-        key,
-        value,
-    ) in artifacts.items():
-        if key not in inputs or not inputs[key]:
-            inputs[key] = value
+
+def load_langchain_object(
+    data_graph: Dict[str, Any], session_id: str
+) -> Tuple[Union[Chain, VectorStore], Dict[str, Any], str]:
+    langchain_object, artifacts = get_build_result(data_graph, session_id)
+    session_id = build_sorted_vertices_with_caching.hash
+    logger.debug("Loaded LangChain object")
 
     if langchain_object is None:
-        # Raise user facing error
         raise ValueError(
             "There was an error loading the langchain_object. Please, check all the nodes and try again."
         )
 
-    # Generate result and thought
+    return langchain_object, artifacts, session_id
+
+
+def process_inputs(inputs: Optional[dict], artifacts: Dict[str, Any]) -> dict:
+    if inputs is None:
+        inputs = {}
+
+    for key, value in artifacts.items():
+        if key not in inputs or not inputs[key]:
+            inputs[key] = value
+
+    return inputs
+
+
+def generate_result(langchain_object: Union[Chain, VectorStore], inputs: dict):
     if isinstance(langchain_object, Chain):
         if inputs is None:
             raise ValueError("Inputs must be provided for a Chain")
@@ -130,9 +146,28 @@ def process_graph_cached(
         raise ValueError(
             f"Unknown langchain_object type: {type(langchain_object).__name__}"
         )
+
     return result
 
 
+def process_graph_cached(
+    data_graph: Dict[str, Any],
+    inputs: Optional[dict] = None,
+    clear_cache=False,
+    session_id=None,
+) -> Tuple[Any, str]:
+    clear_caches_if_needed(clear_cache)
+    # If session_id is provided, load the langchain_object from the session
+    # else build the graph and return the result and the new session_id
+    langchain_object, artifacts, session_id = load_langchain_object(
+        data_graph, session_id
+    )
+    processed_inputs = process_inputs(inputs, artifacts)
+    result = generate_result(langchain_object, processed_inputs)
+
+    return result, session_id
+
+
 def load_flow_from_json(
     flow: Union[Path, str, dict], tweaks: Optional[dict] = None, build=True
 ):

src/backend/langflow/services/cache/utils.py

@@ -30,6 +30,7 @@ def create_cache_folder(func):
 
 def memoize_dict(maxsize=128):
     cache = OrderedDict()
+    hash_to_key = {}  # Mapping from hash to cache key
 
     def decorator(func):
         @functools.wraps(func)
@@ -39,16 +40,29 @@ def memoize_dict(maxsize=128):
             if key not in cache:
                 result = func(*args, **kwargs)
                 cache[key] = result
+                hash_to_key[hashed] = key  # Store the mapping
                 if len(cache) > maxsize:
-                    cache.popitem(last=False)
+                    oldest_key = next(iter(cache))
+                    oldest_hash = oldest_key[1]
+                    del cache[oldest_key]
+                    del hash_to_key[oldest_hash]
             else:
                 result = cache[key]
+
+            wrapper.session_id = hashed  # Store hash in the wrapper
             return result
 
         def clear_cache():
             cache.clear()
+            hash_to_key.clear()
+
+        def get_result_by_session_id(session_id):
+            key = hash_to_key.get(session_id)
+            return cache.get(key) if key is not None else None
 
         wrapper.clear_cache = clear_cache  # type: ignore
+        wrapper.get_result_by_session_id = get_result_by_session_id  # type: ignore
+        wrapper.hash = None
         wrapper.cache = cache  # type: ignore
         return wrapper