升级新库

bbit_ai/app/llm/memLLM.py

@@ -1,12 +1,12 @@
-from langchain.prompts import PromptTemplate
-from config.llm import llm,llmThink
+from langchain_core.prompts import PromptTemplate
+
 import db.milvus as milvus
 import db.postgres as pg
-import json
+from config.llm import llmThink
 
 memPathPrompt = PromptTemplate(
     input_variables=["ai_role", "CHAT_RECORD"],
-    template = """
+    template="""
 你是一个记忆筛选器，负责判断最近对话的信息中，用户的回复内容是否对业务具有长期价值或潜在价值，或者可以帮助形成用户画像。
 首先，请仔细阅读以下关于你业务的描述： 
 <ai_role> 
@@ -31,12 +31,12 @@ no：用户最新回复价值有限或几乎不会在未来业务中使用。
 
 回复不要带任何标点符号以及空格、换行符。
 请给出你的判断结果：
-    """
+    """,
 )
 memPathChain = memPathPrompt | llmThink
 memPrompt = PromptTemplate(
     input_variables=["CHAT_RECORD"],
-    template = """
+    template="""
 你的任务是对给定的聊天记录进行关键信息的记忆总结。请仔细阅读以下聊天记录，并按照要求进行总结：
 <聊天记录>
 {CHAT_RECORD}
@@ -48,13 +48,15 @@ memPrompt = PromptTemplate(
 4. 总结内容应包含时间，并确保时间是准确的。
 5. 你需要针对你的业务场景{ai_role}，展开对用户最后回复的总结。
 请生成你的总结，以用户、时间开头：
-    """
+    """,
 )
 memChain = memPrompt | llmThink
-def take_memory(ai_id:str,sessionId: str,user_id:str, max_retry=3):
+
+
+def take_memory(ai_id: str, sessionId: str, user_id: str, max_retry=3):
     """根据用户输入选择数据来源"""
-    history = pg.get_history_with_time(sessionId,10)
-    print("获取的历史记录:",history)
+    history = pg.get_history_with_time(sessionId, 10)
+    print("获取的历史记录:", history)
     ai_service = pg.get_description(ai_id)
     if ai_service == "":
         # AI描述没有描述，则取业务字段
@@ -66,17 +68,29 @@ def take_memory(ai_id:str,sessionId: str,user_id:str, max_retry=3):
         else:
             ai_service = json["业务"]
     print("获取的描述是：", ai_service)
-    choice = memPathChain.invoke({
-        "ai_role": ai_service,
-        "CHAT_RECORD": history,
-    }).content.strip().lower()
+    choice = (
+        memPathChain.invoke(
+            {
+                "ai_role": ai_service,
+                "CHAT_RECORD": history,
+            }
+        )
+        .content.strip()
+        .lower()
+    )
     print("记忆判断器判断的结果是：", choice)
     if choice == "yes":
         # 对对话进行总结
-        memory = memChain.invoke({
-            "CHAT_RECORD": history,
-            "ai_role": ai_service,
-        }).content.strip().lower()
+        memory = (
+            memChain.invoke(
+                {
+                    "CHAT_RECORD": history,
+                    "ai_role": ai_service,
+                }
+            )
+            .content.strip()
+            .lower()
+        )
         print("记忆生成结果是：", memory)
-        milvus.add_memory(mem = memory,user_id = user_id, is_active = True, ai_id = ai_id)
-        return
+        milvus.add_memory(mem=memory, user_id=user_id, is_active=True, ai_id=ai_id)
+        return