From 020488e97b24950bed6ac0423c9078eca11f9f8b Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 21 Nov 2023 14:49:47 -0300
Subject: [PATCH] Fix typos and update method names

---
 src/backend/langflow/api/v1/callback.py       | 30 ++++++++-----------
 .../langflow/interface/custom/code_parser.py  |  6 ++--
 src/backend/langflow/services/chat/service.py | 18 +++++------
 .../langflow/services/settings/base.py        |  2 +-
 src/backend/langflow/settings.py              |  7 +++--
 src/backend/langflow/template/field/base.py   |  2 +-
 src/backend/langflow/worker.py                | 10 +++----
 tests/conftest.py                             |  8 ++---
 tests/test_database.py                        | 24 +++++++--------
 tests/test_user.py                            | 16 +++++-----
 10 files changed, 60 insertions(+), 63 deletions(-)

diff --git a/src/backend/langflow/api/v1/callback.py b/src/backend/langflow/api/v1/callback.py
index da9005bd0..a838a8750 100644
--- a/src/backend/langflow/api/v1/callback.py
+++ b/src/backend/langflow/api/v1/callback.py
@@ -1,19 +1,15 @@
 import asyncio
+from typing import Any, Dict, List, Optional
 from uuid import UUID
 
 from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
-
-from langflow.api.v1.schemas import ChatResponse, PromptResponse
-
-
-from typing import Any, Dict, List, Optional
-from langflow.services.deps import get_chat_service
-
-
-from langflow.utils.util import remove_ansi_escape_codes
 from langchain.schema import AgentAction, AgentFinish
 from loguru import logger
 
+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
+
 
 # https://github.com/hwchase17/chat-langchain/blob/master/callback.py
 class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
@@ -26,7 +22,7 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
 
     async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
         resp = ChatResponse(message=token, type="stream", intermediate_steps="")
-        await self.websocket.send_json(resp.dict())
+        await self.websocket.send_json(resp.model_dump())
 
     async def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
         """Run when tool starts running."""
@@ -35,7 +31,7 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
             type="stream",
             intermediate_steps=f"Tool input: {input_str}",
         )
-        await self.websocket.send_json(resp.dict())
+        await self.websocket.send_json(resp.model_dump())
 
     async def on_tool_end(self, output: str, **kwargs: Any) -> Any:
         """Run when tool ends running."""
@@ -66,7 +62,7 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
         try:
             # This is to emulate the stream of tokens
             for resp in resps:
-                await self.websocket.send_json(resp.dict())
+                await self.websocket.send_json(resp.model_dump())
         except Exception as exc:
             logger.error(f"Error sending response: {exc}")
 
@@ -92,7 +88,7 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
             resp = PromptResponse(
                 prompt=text,
             )
-            await self.websocket.send_json(resp.dict())
+            await self.websocket.send_json(resp.model_dump())
             self.chat_service.chat_history.add_message(self.client_id, resp)
 
     async def on_agent_action(self, action: AgentAction, **kwargs: Any):
@@ -103,10 +99,10 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
             logs = log.split("\n")
             for log in logs:
                 resp = ChatResponse(message="", type="stream", intermediate_steps=log)
-                await self.websocket.send_json(resp.dict())
+                await self.websocket.send_json(resp.model_dump())
         else:
             resp = ChatResponse(message="", type="stream", intermediate_steps=log)
-            await self.websocket.send_json(resp.dict())
+            await self.websocket.send_json(resp.model_dump())
 
     async def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
         """Run on agent end."""
@@ -115,7 +111,7 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler):
             type="stream",
             intermediate_steps=finish.log,
         )
-        await self.websocket.send_json(resp.dict())
+        await self.websocket.send_json(resp.model_dump())
 
 
 class StreamingLLMCallbackHandler(BaseCallbackHandler):
@@ -130,5 +126,5 @@ class StreamingLLMCallbackHandler(BaseCallbackHandler):
         resp = ChatResponse(message=token, type="stream", intermediate_steps="")
 
         loop = asyncio.get_event_loop()
-        coroutine = self.websocket.send_json(resp.dict())
+        coroutine = self.websocket.send_json(resp.model_dump())
         asyncio.run_coroutine_threadsafe(coroutine, loop)
diff --git a/src/backend/langflow/interface/custom/code_parser.py b/src/backend/langflow/interface/custom/code_parser.py
index 5202ed507..ce869d263 100644
--- a/src/backend/langflow/interface/custom/code_parser.py
+++ b/src/backend/langflow/interface/custom/code_parser.py
@@ -1,8 +1,8 @@
 import ast
 import inspect
 import traceback
+from typing import Any, Dict, List, Type, Union
 
-from typing import Dict, Any, List, Type, Union
 from fastapi import HTTPException
 from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails
 
@@ -104,7 +104,7 @@ class CodeParser:
         func.args = self.parse_function_args(node)
         func.body = self.parse_function_body(node)
 
-        return func.dict()
+        return func.model_dump()
 
     def parse_function_args(self, node: ast.FunctionDef) -> List[Dict[str, Any]]:
         """
@@ -233,7 +233,7 @@ class CodeParser:
                 else:
                     class_details.methods.append(method)
 
-        self.data["classes"].append(class_details.dict())
+        self.data["classes"].append(class_details.model_dump())
 
     def parse_global_vars(self, node: ast.Assign) -> None:
         """
diff --git a/src/backend/langflow/services/chat/service.py b/src/backend/langflow/services/chat/service.py
index 58d7cf384..1287f08ba 100644
--- a/src/backend/langflow/services/chat/service.py
+++ b/src/backend/langflow/services/chat/service.py
@@ -1,20 +1,20 @@
-from collections import defaultdict
+import asyncio
 import uuid
+from collections import defaultdict
+from typing import Any, Dict, List
+
+import orjson
 from fastapi import WebSocket, status
-from starlette.websockets import WebSocketState
 from langflow.api.v1.schemas import ChatMessage, ChatResponse, FileResponse
 from langflow.interface.utils import pil_to_base64
+from langflow.services import ServiceType, service_manager
 from langflow.services.base import Service
 from langflow.services.chat.cache import Subject
 from langflow.services.chat.utils import process_graph
 from loguru import logger
+from starlette.websockets import WebSocketState
 
 from .cache import cache_service
-import asyncio
-from typing import Any, Dict, List
-
-from langflow.services import service_manager, ServiceType
-import orjson
 
 
 class ChatHistory(Subject):
@@ -104,7 +104,7 @@ class ChatService(Service):
 
     async def send_json(self, client_id: str, message: ChatMessage):
         websocket = self.active_connections[client_id]
-        await websocket.send_json(message.dict())
+        await websocket.send_json(message.model_dump())
 
     async def close_connection(self, client_id: str, code: int, reason: str):
         if websocket := self.active_connections[client_id]:
@@ -191,7 +191,7 @@ class ChatService(Service):
         try:
             chat_history = self.chat_history.get_history(client_id)
             # iterate and make BaseModel into dict
-            chat_history = [chat.dict() for chat in chat_history]
+            chat_history = [chat.model_dump() for chat in chat_history]
             await websocket.send_json(chat_history)
 
             while True:
diff --git a/src/backend/langflow/services/settings/base.py b/src/backend/langflow/services/settings/base.py
index 53c6ef947..63c9de38d 100644
--- a/src/backend/langflow/services/settings/base.py
+++ b/src/backend/langflow/services/settings/base.py
@@ -205,7 +205,7 @@ class Settings(BaseSettings):
 
 def save_settings_to_yaml(settings: Settings, file_path: str):
     with open(file_path, "w") as f:
-        settings_dict = settings.dict()
+        settings_dict = settings.model_dump()
         yaml.dump(settings_dict, f)
 
 
diff --git a/src/backend/langflow/settings.py b/src/backend/langflow/settings.py
index 3b06fd511..625bd56e5 100644
--- a/src/backend/langflow/settings.py
+++ b/src/backend/langflow/settings.py
@@ -1,12 +1,13 @@
 import contextlib
 import json
 import os
-from typing import Optional, List
 from pathlib import Path
+from typing import List, Optional
 
 import yaml
-from pydantic import validator, model_validator
+from pydantic import model_validator, validator
 from pydantic_settings import BaseSettings
+
 from langflow.utils.logger import logger
 
 BASE_COMPONENTS_PATH = str(Path(__file__).parent / "components")
@@ -141,7 +142,7 @@ class Settings(BaseSettings):
 
 def save_settings_to_yaml(settings: Settings, file_path: str):
     with open(file_path, "w") as f:
-        settings_dict = settings.dict()
+        settings_dict = settings.model_dump()
         yaml.dump(settings_dict, f)
 
 
diff --git a/src/backend/langflow/template/field/base.py b/src/backend/langflow/template/field/base.py
index 31c68d094..e596f21be 100644
--- a/src/backend/langflow/template/field/base.py
+++ b/src/backend/langflow/template/field/base.py
@@ -60,7 +60,7 @@ class TemplateFieldCreator(BaseModel, ABC):
     """Additional information about the field to be shown in the tooltip. Defaults to an empty string."""
 
     def to_dict(self):
-        result = self.dict()
+        result = self.model_dump()
         # Remove key if it is None
         for key in list(result.keys()):
             if result[key] is None or result[key] == []:
diff --git a/src/backend/langflow/worker.py b/src/backend/langflow/worker.py
index e466792cf..8f2abcb43 100644
--- a/src/backend/langflow/worker.py
+++ b/src/backend/langflow/worker.py
@@ -1,15 +1,15 @@
-from langflow.core.celery_app import celery_app
-from typing import Any, Dict, Optional
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from celery.exceptions import SoftTimeLimitExceeded  # type: ignore
+
+from langflow.core.celery_app import celery_app
 from langflow.processing.process import (
     Result,
     generate_result,
     process_inputs,
 )
-from langflow.services.manager import initialize_session_service
 from langflow.services.deps import get_session_service
+from langflow.services.manager import initialize_session_service
 
 if TYPE_CHECKING:
     from langflow.graph.vertex.base import Vertex
@@ -55,4 +55,4 @@ def process_graph_cached_task(
     # we need to update the cache with the updated langchain_object
     session_service.update_session(session_id, (graph, artifacts))
 
-    return Result(result=result, session_id=session_id).dict()
+    return Result(result=result, session_id=session_id).model_dump()
diff --git a/tests/conftest.py b/tests/conftest.py
index 238e7cdab..0acc3116f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -224,7 +224,7 @@ def test_user(client):
         username="testuser",
         password="testpassword",
     )
-    response = client.post("/api/v1/users", json=user_data.dict())
+    response = client.post("/api/v1/users", json=user_data.model_dump())
     assert response.status_code == 201
     return response.json()
 
@@ -269,7 +269,7 @@ def flow(client, json_flow: str, active_user):
         user_id=active_user.id,
         description="description",
     )
-    flow = Flow(**flow_data.dict())
+    flow = Flow(**flow_data.model_dump())
     with session_getter(get_db_service()) as session:
         session.add(flow)
         session.commit()
@@ -283,7 +283,7 @@ def added_flow(client, json_flow_with_prompt_and_history, logged_in_headers):
     flow = orjson.loads(json_flow_with_prompt_and_history)
     data = flow["data"]
     flow = FlowCreate(name="Basic Chat", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
@@ -295,7 +295,7 @@ def added_vector_store(client, json_vector_store, logged_in_headers):
     vector_store = orjson.loads(json_vector_store)
     data = vector_store["data"]
     vector_store = FlowCreate(name="Vector Store", description="description", data=data)
-    response = client.post("api/v1/flows/", json=vector_store.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=vector_store.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
     assert response.json()["name"] == vector_store.name
     assert response.json()["data"] == vector_store.data
diff --git a/tests/test_database.py b/tests/test_database.py
index 7a795f33a..4a1219799 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -29,7 +29,7 @@ def test_create_flow(client: TestClient, json_flow: str, active_user, logged_in_
     flow = orjson.loads(json_flow)
     data = flow["data"]
     flow = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
@@ -45,13 +45,13 @@ def test_read_flows(client: TestClient, json_flow: str, active_user, logged_in_h
     flow_data = orjson.loads(json_flow)
     data = flow_data["data"]
     flow = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
 
     flow = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
@@ -65,7 +65,7 @@ def test_read_flow(client: TestClient, json_flow: str, active_user, logged_in_he
     flow = orjson.loads(json_flow)
     data = flow["data"]
     flow = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     flow_id = response.json()["id"]  # flow_id should be a UUID but is a string
     # turn it into a UUID
     flow_id = UUID(flow_id)
@@ -81,7 +81,7 @@ def test_update_flow(client: TestClient, json_flow: str, active_user, logged_in_
     data = flow["data"]
 
     flow = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
 
     flow_id = response.json()["id"]
     updated_flow = FlowUpdate(
@@ -89,7 +89,7 @@ def test_update_flow(client: TestClient, json_flow: str, active_user, logged_in_
         description="updated description",
         data=data,
     )
-    response = client.patch(f"api/v1/flows/{flow_id}", json=updated_flow.dict(), headers=logged_in_headers)
+    response = client.patch(f"api/v1/flows/{flow_id}", json=updated_flow.model_dump(), headers=logged_in_headers)
 
     assert response.status_code == 200
     assert response.json()["name"] == updated_flow.name
@@ -101,7 +101,7 @@ def test_delete_flow(client: TestClient, json_flow: str, active_user, logged_in_
     flow = orjson.loads(json_flow)
     data = flow["data"]
     flow = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     flow_id = response.json()["id"]
     response = client.delete(f"api/v1/flows/{flow_id}", headers=logged_in_headers)
     assert response.status_code == 200
@@ -119,7 +119,7 @@ def test_create_flows(client: TestClient, session: Session, json_flow: str, logg
         ]
     )
     # Make request to endpoint
-    response = client.post("api/v1/flows/batch/", json=flow_list.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/batch/", json=flow_list.model_dump(), headers=logged_in_headers)
     # Check response status code
     assert response.status_code == 201
     # Check response data
@@ -143,7 +143,7 @@ def test_upload_file(client: TestClient, session: Session, json_flow: str, logge
             FlowCreate(name="Flow 2", description="description", data=data),
         ]
     )
-    file_contents = orjson_dumps(flow_list.dict())
+    file_contents = orjson_dumps(flow_list.model_dump())
     response = client.post(
         "api/v1/flows/upload/",
         files={"file": ("examples.json", file_contents, "application/json")},
@@ -216,11 +216,11 @@ def test_update_flow_idempotency(client: TestClient, json_flow: str, active_user
     flow_data = orjson.loads(json_flow)
     data = flow_data["data"]
     flow_data = FlowCreate(name="Test Flow", description="description", data=data)
-    response = client.post("api/v1/flows/", json=flow_data.dict(), headers=logged_in_headers)
+    response = client.post("api/v1/flows/", json=flow_data.model_dump(), headers=logged_in_headers)
     flow_id = response.json()["id"]
     updated_flow = FlowCreate(name="Updated Flow", description="description", data=data)
-    response1 = client.put(f"api/v1/flows/{flow_id}", json=updated_flow.dict(), headers=logged_in_headers)
-    response2 = client.put(f"api/v1/flows/{flow_id}", json=updated_flow.dict(), headers=logged_in_headers)
+    response1 = client.put(f"api/v1/flows/{flow_id}", json=updated_flow.model_dump(), headers=logged_in_headers)
+    response2 = client.put(f"api/v1/flows/{flow_id}", json=updated_flow.model_dump(), headers=logged_in_headers)
     assert response1.json() == response2.json()
 
 
diff --git a/tests/test_user.py b/tests/test_user.py
index 2933c5b18..e93b594cc 100644
--- a/tests/test_user.py
+++ b/tests/test_user.py
@@ -1,11 +1,11 @@
 from datetime import datetime
-from langflow.services.auth.utils import create_super_user, get_password_hash
 
+import pytest
+from langflow.services.auth.utils import create_super_user, get_password_hash
+from langflow.services.database.models.user import UserUpdate
 from langflow.services.database.models.user.user import User
 from langflow.services.database.utils import session_getter
 from langflow.services.deps import get_db_service, get_settings_service
-import pytest
-from langflow.services.database.models.user import UserUpdate
 
 
 @pytest.fixture
@@ -89,7 +89,7 @@ def test_data_consistency_after_update(client, active_user, logged_in_headers, s
     user_id = active_user.id
     update_data = UserUpdate(is_active=False)
 
-    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.dict(), headers=super_user_headers)
+    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.model_dump(), headers=super_user_headers)
     assert response.status_code == 200, response.json()
 
     # Fetch the updated user from the database
@@ -163,13 +163,13 @@ def test_patch_user(client, active_user, logged_in_headers):
         username="newname",
     )
 
-    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.dict(), headers=logged_in_headers)
+    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.model_dump(), headers=logged_in_headers)
     assert response.status_code == 200, response.json()
     update_data = UserUpdate(
         profile_image="new_image",
     )
 
-    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.dict(), headers=logged_in_headers)
+    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.model_dump(), headers=logged_in_headers)
     assert response.status_code == 200, response.json()
 
 
@@ -181,7 +181,7 @@ def test_patch_reset_password(client, active_user, logged_in_headers):
 
     response = client.patch(
         f"/api/v1/users/{user_id}/reset-password",
-        json=update_data.dict(),
+        json=update_data.model_dump(),
         headers=logged_in_headers,
     )
     assert response.status_code == 200, response.json()
@@ -197,7 +197,7 @@ def test_patch_user_wrong_id(client, active_user, logged_in_headers):
         username="newname",
     )
 
-    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.dict(), headers=logged_in_headers)
+    response = client.patch(f"/api/v1/users/{user_id}", json=update_data.model_dump(), headers=logged_in_headers)
     assert response.status_code == 422, response.json()
     assert response.json() == {
         "detail": [