dev新增短期记忆功能 (#47)

* dev新增短期记忆功能 * dev新增短期记忆功能 * dev新增短期记忆功能 * dev新增短期记忆功能 * dev新增短期记忆功能 * dev新增短期记忆功能 * dev新增短期记忆功能
2026-01-07 16:36:11 +08:00
parent 5fe8043ff8
commit bcb3d587a1
9 changed files with 765 additions and 45 deletions
--- a/api/app/core/agent/langchain_agent.py
+++ b/api/app/core/agent/langchain_agent.py
@@ -7,13 +7,20 @@ LangChain Agent 封装
 - 支持流式输出
 - 使用 RedBearLLM 支持多提供商
 """
+import os
 import time
 from typing import Any, AsyncGenerator, Dict, List, Optional, Sequence

+
+from app.db import get_db
 from app.core.logging_config import get_business_logger
 from app.core.memory.agent.utils.redis_tool import store
 from app.core.models import RedBearLLM, RedBearModelConfig
 from app.models.models_model import ModelType
+from app.repositories.memory_short_repository import LongTermMemoryRepository
+from app.services.memory_agent_service import (
+    get_end_user_connected_config,
+)
 from app.services.memory_konwledges_server import write_rag
 from app.services.task_service import get_task_memory_write_result
 from app.tasks import write_message_task
@@ -96,7 +103,8 @@ class LangChainAgent:
                "temperature": temperature,
                "streaming": streaming,
                "tool_count": len(self.tools),
-                "tool_names": [tool.name for tool in self.tools] if self.tools else []
+                "tool_names": [tool.name for tool in self.tools] if self.tools else [],
+                "tool_count": len(self.tools)
            }
        )

@@ -137,11 +145,8 @@ class LangChainAgent:
        messages.append(HumanMessage(content=user_content))

        return messages
-
    async def term_memory_save(self,messages,end_user_end,aimessages):
-        """
-        短长期存储redis，为不影响正常使用6句一段话，存储用户名加一个前缀，当数据存够6条返回给neo4j
-        """
+        '''短长期存储redis，为不影响正常使用6句一段话，存储用户名加一个前缀，当数据存够6条返回给neo4j'''
        end_user_end=f"Term_{end_user_end}"
        print(messages)
        print(aimessages)
@@ -155,17 +160,18 @@ class LangChainAgent:
        store.delete_duplicate_sessions()
        # logger.info(f'Redis_Agent:{end_user_end};{session_id}')
        return session_id
-
    async def term_memory_redis_read(self,end_user_end):
        end_user_end = f"Term_{end_user_end}"
        history = store.find_user_apply_group(end_user_end, end_user_end, end_user_end)
        # logger.info(f'Redis_Agent:{end_user_end};{history}')
        messagss_list=[]
+        retrieved_content=[]
        for messages in history:
            query = messages.get("Query")
            aimessages = messages.get("Answer")
            messagss_list.append(f'用户:{query}。AI回复:{aimessages}')
-        return messagss_list
+            retrieved_content.append({query: aimessages})
+        return messagss_list,retrieved_content


    async def write(self,storage_type,end_user_id,message,user_rag_memory_id,actual_end_user_id,content,actual_config_id):
@@ -205,7 +211,6 @@ class LangChainAgent:
        # If config_id is None, try to get from end_user's connected config
        if actual_config_id is None and end_user_id:
            try:
-                from app.db import get_db
                from app.services.memory_agent_service import (
                    get_end_user_connected_config,
                )
@@ -223,11 +228,26 @@ class LangChainAgent:
        logger.info(f'写入类型{storage_type,str(end_user_id), message, str(user_rag_memory_id)}')
        print(f'写入类型{storage_type,str(end_user_id), message, str(user_rag_memory_id)}')

-        history_term_memory=await self.term_memory_redis_read(end_user_id)
+        history_term_memory_result = await self.term_memory_redis_read(end_user_id)
+        history_term_memory = history_term_memory_result[0]
+        db_for_memory = next(get_db())
        if memory_flag:
            if len(history_term_memory)>=4 and storage_type != "rag":
-                history_term_memory=';'.join(history_term_memory)
-                logger.info(f'写入短长期：{storage_type, str(end_user_id), history_term_memory, str(user_rag_memory_id)}')
+                history_term_memory = ';'.join(history_term_memory)
+                retrieved_content = history_term_memory_result[1]
+                print(retrieved_content)
+                # 为长期记忆操作获取新的数据库连接
+                try:
+                    repo = LongTermMemoryRepository(db_for_memory)
+                    repo.upsert(end_user_id, retrieved_content)
+                    logger.info(
+                        f'写入短长期：{storage_type, str(end_user_id), history_term_memory, str(user_rag_memory_id)}')
+                except Exception as e:
+                    logger.error(f"Failed to write to LongTermMemory: {e}")
+                    raise
+                finally:
+                    db_for_memory.close()
+
                await self.write(storage_type,end_user_id,history_term_memory,user_rag_memory_id,actual_end_user_id,history_term_memory,actual_config_id)
            await self.write(storage_type,end_user_id,message,user_rag_memory_id,actual_end_user_id,message,actual_config_id)
        try:
@@ -316,10 +336,6 @@ class LangChainAgent:
        # If config_id is None, try to get from end_user's connected config
        if actual_config_id is None and end_user_id:
            try:
-                from app.db import get_db
-                from app.services.memory_agent_service import (
-                    get_end_user_connected_config,
-                )
                db = next(get_db())
                try:
                    connected_config = get_end_user_connected_config(end_user_id, db)
@@ -331,14 +347,24 @@ class LangChainAgent:
            except Exception as e:
                logger.warning(f"Failed to get db session: {e}")

-        history_term_memory = await self.term_memory_redis_read(end_user_id)
+        history_term_memory_result = await self.term_memory_redis_read(end_user_id)
+        history_term_memory = history_term_memory_result[0]
        if memory_flag:
            if len(history_term_memory) >= 4 and storage_type != "rag":
                history_term_memory = ';'.join(history_term_memory)
-                logger.info(
-                    f'写入短长期：{storage_type, str(end_user_id), history_term_memory, str(user_rag_memory_id)}')
-                await self.write(storage_type, end_user_id, history_term_memory, user_rag_memory_id, end_user_id,
-                                 history_term_memory, actual_config_id)
+                retrieved_content = history_term_memory_result[1]
+                db_for_memory = next(get_db())
+                try:
+                    repo = LongTermMemoryRepository(db_for_memory)
+                    repo.upsert(end_user_id, retrieved_content)
+                    logger.info(
+                        f'写入短长期：{storage_type, str(end_user_id), history_term_memory, str(user_rag_memory_id)}')
+                    await self.write(storage_type, end_user_id, history_term_memory, user_rag_memory_id, end_user_id,
+                                     history_term_memory, actual_config_id)
+                except Exception as e:
+                    logger.error(f"Failed to write to long term memory: {e}")
+                finally:
+                    db_for_memory.close()

            await self.write(storage_type, end_user_id, message, user_rag_memory_id, end_user_id, message, actual_config_id)
        try: