Skip to content

Commit

Permalink
Fix some blocking usages of messages memory
Browse files Browse the repository at this point in the history
  • Loading branch information
cbornet committed Nov 11, 2024
1 parent 69a1c16 commit 366e45b
Show file tree
Hide file tree
Showing 4 changed files with 11 additions and 5 deletions.
3 changes: 2 additions & 1 deletion src/backend/base/langflow/base/agents/events.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# Add helper functions for each event type
import asyncio
from collections.abc import AsyncIterator
from time import perf_counter
from typing import Any, Protocol
Expand Down Expand Up @@ -235,7 +236,7 @@ async def process_agent_events(
agent_message.properties.icon = "Bot"
agent_message.properties.state = "partial"
# Store the initial message
agent_message = send_message_method(message=agent_message)
agent_message = asyncio.to_thread(send_message_method, message=agent_message)
try:
# Create a mapping of run_ids to tool contents
tool_blocks_map: dict[str, ToolContent] = {}
Expand Down
4 changes: 3 additions & 1 deletion src/backend/base/langflow/components/agents/agent.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import asyncio

from langchain_core.tools import StructuredTool

from langflow.base.agents.agent import LCToolsAgentComponent
Expand Down Expand Up @@ -58,7 +60,7 @@ async def message_response(self) -> Message:
if llm_model is None:
msg = "No language model selected"
raise ValueError(msg)
self.chat_history = self.get_memory_data()
self.chat_history = await asyncio.to_thread(self.get_memory_data)

if self.add_current_date_tool:
if not isinstance(self.tools, list): # type: ignore[has-type]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -797,15 +797,17 @@ async def build_results(self):
return await self._build_with_tracing()
return await self._build_without_tracing()
except StreamingError as e:
self.send_error(
await asyncio.to_thread(
self.send_error,
exception=e.cause,
session_id=self.graph.session_id,
trace_name=getattr(self, "trace_name", None),
source=e.source,
)
raise e.cause # noqa: B904
except Exception as e:
self.send_error(
await asyncio.to_thread(
self.send_error,
exception=e,
session_id=self.graph.session_id,
source=Source(id=self._id, display_name=self.display_name, source=self.display_name),
Expand Down
3 changes: 2 additions & 1 deletion src/backend/base/langflow/graph/vertex/types.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import asyncio
import contextlib
import json
from collections.abc import AsyncIterator, Generator, Iterator
Expand Down Expand Up @@ -434,7 +435,7 @@ async def stream(self):
and hasattr(self.custom_component, "should_store_message")
and hasattr(self.custom_component, "store_message")
):
self.custom_component.store_message(message)
await asyncio.to_thread(self.custom_component.store_message, message)
log_vertex_build(
flow_id=self.graph.flow_id,
vertex_id=self.id,
Expand Down

0 comments on commit 366e45b

Please sign in to comment.