Merge pull request #40 from chatchat-space/dev

yuehua-s · web-flow · commit 6e72333211e4 · 2024-11-11T19:31:18.000+08:00
Dev
diff --git a/libs/chatchat-server/chatchat/server/utils.py b/libs/chatchat-server/chatchat/server/utils.py
@@ -25,7 +25,7 @@
 from langchain_openai.chat_models import ChatOpenAI
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
-# from langgraph.checkpoint.postgres import PostgresSaver
+from langgraph.checkpoint.postgres import PostgresSaver
 from memoization import cached, CachingAlgorithmFlag
 
 from chatchat.settings import Settings, XF_MODELS_TYPES
@@ -774,28 +774,17 @@ def get_tool_config(name: str = None) -> Dict:
         return Settings.tool_settings.model_dump().get(name, {})
 
 
-def get_checkpoint(memory_type: Optional[Literal["memory", "sqlite", "postgres"]] = None) -> (Union)[
-    MemorySaver,
-    AsyncSqliteSaver,
-    # PostgresSaver
-]:
+def get_checkpointer(memory_type: Optional[Literal["memory", "sqlite", "postgres"]] = None) -> (Union)[MemorySaver, AsyncSqliteSaver]:
     """
-    获取 graph 的 memory
+    获取 graph 的 checkpointer(MemorySaver 和 AsyncSqliteSaver)
     """
-    import sqlalchemy as sa
-
     if memory_type is None:
         memory_type = Settings.tool_settings.GRAPH_MEMORY_TYPE
 
     if memory_type == "memory":
         return MemorySaver()
     elif memory_type == "sqlite":
         return AsyncSqliteSaver.from_conn_string(Settings.basic_settings.SQLITE_GRAPH_DATABASE_URI)
-    # elif memory_type == "postgres":
-    #     import sqlalchemy as sa
-    #     engine = sa.create_engine(Settings.basic_settings.SQLALCHEMY_DATABASE_URI)
-    #     conn = engine.connect().connection
-    #     return PostgresSaver(conn)
 
     raise ValueError("Invalid memory_type provided. Must be 'memory', 'sqlite', or 'postgres'.")
 
diff --git a/libs/chatchat-server/chatchat/settings.py b/libs/chatchat-server/chatchat/settings.py
@@ -101,6 +101,18 @@ def BASE_TEMP_DIR(self) -> Path:
     SQLITE_GRAPH_DATABASE_URI: str = str(CHATCHAT_ROOT / "data/graph.db")
     """工作流 SQLITE CHECKPOINT 数据库连接URI"""
 
+    POSTGRESQL_GRAPH_DATABASE_URI: str = "postgresql://username:password@localhost:5442/langgraph_chatchat"
+    """工作流 POSTGRESQL CHECKPOINT 数据库连接URI"""
+
+    POSTGRESQL_GRAPH_CONNECTION_POOLS_MAX_SIZE: int = 20
+    """工作流 POSTGRESQL CHECKPOINT 数据库连接池最大连接数限制"""
+
+    POSTGRESQL_GRAPH_CONNECTION_POOLS_KWARGS: dict = {
+        "autocommit": True,
+        "prepare_threshold": 0,
+    }
+    """工作流 POSTGRESQL CHECKPOINT 数据库连接池关键字参数配置"""
+
     OPEN_CROSS_DOMAIN: bool = False
     """API 是否开启跨域"""
 
diff --git a/libs/chatchat-server/chatchat/webui_pages/graph_agent/graph.py b/libs/chatchat-server/chatchat/webui_pages/graph_agent/graph.py
@@ -13,7 +13,7 @@
     get_tool,
     list_tools,
     create_agent_models,
-    get_checkpoint
+    get_checkpointer
 )
 
 logger = build_logger()
@@ -153,7 +153,7 @@ async def create_graph(
 ):
     if st.session_state["checkpoint_type"] == "memory":
         if "memory" not in st.session_state:
-            st.session_state["memory"] = get_checkpoint()
+            st.session_state["memory"] = get_checkpointer()
         checkpoint = st.session_state["memory"]
         graph_class = graph_class(llm=graph_llm,
                                   tools=graph_tools,
@@ -164,7 +164,7 @@ async def create_graph(
             raise ValueError(f"Graph '{graph_class}' is not registered.")
         await process_graph(graph_class=graph_class, graph=graph, graph_input=graph_input, graph_config=graph_config)
     elif st.session_state["checkpoint_type"] == "sqlite":
-        checkpoint_class = get_checkpoint()
+        checkpoint_class = get_checkpointer()
         async with checkpoint_class as checkpoint:
             graph_class = graph_class(llm=graph_llm,
                                       tools=graph_tools,
@@ -174,6 +174,25 @@ async def create_graph(
             if not graph:
                 raise ValueError(f"Graph '{graph_class}' is not registered.")
             await process_graph(graph_class=graph_class, graph=graph, graph_input=graph_input, graph_config=graph_config)
+    elif st.session_state["checkpoint_type"] == "postgres":
+        from psycopg_pool import AsyncConnectionPool
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+        async with AsyncConnectionPool(
+            conninfo=Settings.basic_settings.POSTGRESQL_GRAPH_DATABASE_URI,
+            max_size=Settings.basic_settings.POSTGRESQL_GRAPH_CONNECTION_POOLS_MAX_SIZE,
+            kwargs=Settings.basic_settings.POSTGRESQL_GRAPH_CONNECTION_POOLS_KWARGS,
+        ) as pool:
+            checkpoint = AsyncPostgresSaver(pool)
+            # NOTE: you need to call .setup() the first time you're using your checkpointer
+            await checkpoint.setup()
+            graph_class = graph_class(llm=graph_llm,
+                                      tools=graph_tools,
+                                      history_len=graph_history_len,
+                                      checkpoint=checkpoint)
+            graph = graph_class.get_graph()
+            if not graph:
+                raise ValueError(f"Graph '{graph_class}' is not registered.")
+            await process_graph(graph_class=graph_class, graph=graph, graph_input=graph_input, graph_config=graph_config)
 
 
 async def update_state(graph: CompiledStateGraph, graph_config: Dict, update_message: Dict, as_node: str):
diff --git a/libs/chatchat-server/chatchat/webui_pages/graph_rag/rag.py b/libs/chatchat-server/chatchat/webui_pages/graph_rag/rag.py
@@ -13,7 +13,7 @@
     get_tool,
     create_agent_models,
     list_tools,
-    get_checkpoint,
+    get_checkpointer,
 )
 
 logger = build_logger()
@@ -32,7 +32,7 @@ async def create_graph(
 ):
     if st.session_state["checkpoint_type"] == "memory":
         if "memory" not in st.session_state:
-            st.session_state["memory"] = get_checkpoint()
+            st.session_state["memory"] = get_checkpointer()
         checkpoint = st.session_state["memory"]
         graph_class = graph_class(llm=graph_llm,
                                   tools=graph_tools,
@@ -46,7 +46,7 @@ async def create_graph(
             raise ValueError(f"Graph '{graph_class}' is not registered.")
         await process_graph(graph_class=graph_class, graph=graph, graph_input=graph_input, graph_config=graph_config)
     elif st.session_state["checkpoint_type"] == "sqlite":
-        checkpoint_class = get_checkpoint()
+        checkpoint_class = get_checkpointer()
         async with checkpoint_class as checkpoint:
             graph_class = graph_class(llm=graph_llm,
                                       tools=graph_tools,
@@ -59,6 +59,28 @@ async def create_graph(
             if not graph:
                 raise ValueError(f"Graph '{graph_class}' is not registered.")
             await process_graph(graph_class=graph_class, graph=graph, graph_input=graph_input, graph_config=graph_config)
+    elif st.session_state["checkpoint_type"] == "postgres":
+        from psycopg_pool import AsyncConnectionPool
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+        async with AsyncConnectionPool(
+            conninfo=Settings.basic_settings.POSTGRESQL_GRAPH_DATABASE_URI,
+            max_size=Settings.basic_settings.POSTGRESQL_GRAPH_CONNECTION_POOLS_MAX_SIZE,
+            kwargs=Settings.basic_settings.POSTGRESQL_GRAPH_CONNECTION_POOLS_KWARGS,
+        ) as pool:
+            checkpoint = AsyncPostgresSaver(pool)
+            # NOTE: you need to call .setup() the first time you're using your checkpointer
+            await checkpoint.setup()
+            graph_class = graph_class(llm=graph_llm,
+                                      tools=graph_tools,
+                                      history_len=graph_history_len,
+                                      checkpoint=checkpoint,
+                                      knowledge_base=knowledge_base,
+                                      top_k=top_k,
+                                      score_threshold=score_threshold)
+            graph = graph_class.get_graph()
+            if not graph:
+                raise ValueError(f"Graph '{graph_class}' is not registered.")
+            await process_graph(graph_class=graph_class, graph=graph, graph_input=graph_input, graph_config=graph_config)
 
 
 async def graph_rag_page(api: ApiRequest):
diff --git a/libs/chatchat-server/poetry.lock b/libs/chatchat-server/poetry.lock
diff --git a/libs/chatchat-server/pyproject.toml b/libs/chatchat-server/pyproject.toml
@@ -77,6 +77,8 @@ langchain-community = "^0.3.5"
 mysqlclient = "^2.2.5"
 metaphor-python = "^0.1.23"
 xinference = "^0.16.3"
+psycopg-pool = "^3.2.3"
+langgraph-checkpoint-postgres = "^2.0.2"
 [tool.poetry.extras]
 xinference = ["xinference_client"]
 zhipuai = ["zhipuai"]
