加入图片ai功能

app/ai/ai.py

@@ -0,0 +1,131 @@
+
+!pip install -q langgraph-checkpoint-sqlite langchain_google_genai
+
+import sqlite3
+from google.colab import userdata
+from typing import Literal, TypedDict, Annotated
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.messages import SystemMessage, HumanMessage, RemoveMessage, AnyMessage
+from langgraph.graph import StateGraph, START, END
+from langgraph.checkpoint.sqlite import SqliteSaver
+from langgraph.graph.message import add_messages
+from langchain_core.runnables import RunnableConfig
+from typing import Union, List, Dict
+
+
+# --- 1. 状态定义 ---
+class State(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+    summary: str  # 永久存储在数据库中的摘要内容
+
+# --- 2. 核心逻辑 ---
+llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature = 0.7, google_api_key='') #在这里倒入API key
+
+def call_model(state: State, config: RunnableConfig):
+    """对话节点：融合了动态 Prompt 和 长期摘要"""
+
+    # 获取当前 session 特有的 System Prompt（如果没传则使用默认）
+    configurable = config.get("configurable", {})
+    system_base_prompt = configurable.get("system_prompt", "你是一个通用的 AI 助手。")
+
+    # 构造当前上下文
+    summary = state.get("summary", "")
+    if summary:
+        system_base_prompt += f"\n\n<context_summary>\n{summary}\n</context_summary>"
+
+    messages = [SystemMessage(content=system_base_prompt)] + state["messages"]
+    response = llm.invoke(messages)
+    return {"messages": [response]}
+
+def summarize_conversation(state: State):
+    """总结节点：负责更新摘要并清理过期消息"""
+    summary = state.get("summary", "")
+
+    messages_to_summarize = state["messages"][:-1]
+
+    # If there's nothing to summarize yet, just END
+    if not messages_to_summarize:
+        return {"summary": summary}
+
+    system_prompt = (
+        "你是一个记忆管理专家。请更新摘要，合并新旧信息。"
+        "1. 保持简练，仅保留事实（姓名、偏好、核心议题）。"
+        "2. 如果新消息包含对旧信息的修正，请更新它。"
+        "3. 如果对话中包含图片描述，请将图片的关键视觉信息也记录在摘要中"
+    )
+
+    summary_input = f"现有摘要: {summary}\n\n待加入的新信息: {messages_to_summarize}"
+
+    # Invoke model to get new condensed summary
+    response = llm.invoke([
+        SystemMessage(content=system_prompt),
+        HumanMessage(content=summary_input)
+    ])
+
+    # Important: Create RemoveMessage objects for all messages that were summarized
+    delete_messages = [RemoveMessage(id=m.id) for m in messages_to_summarize if m.id]
+
+    return {
+        "summary": response.content,
+        "messages": delete_messages
+    }
+
+def should_continue(state: State) -> Literal["summarize", END]:
+    """如果消息累积超过3条，则去总结节点"""
+    if len(state["messages"]) > 3: #changed
+        return "summarize"
+    return END
+
+# --- 3. 构建图 ---
+db_path = "multi_session_chat.sqlite"
+conn = sqlite3.connect(db_path, check_same_thread=False)
+memory = SqliteSaver(conn)
+
+workflow = StateGraph(State)
+workflow.add_node("chatbot", call_model)
+workflow.add_node("summarize", summarize_conversation)
+
+workflow.add_edge(START, "chatbot")
+workflow.add_conditional_edges("chatbot", should_continue)
+workflow.add_edge("summarize", END)
+
+app = workflow.compile(checkpointer=memory)
+
+def chat(thread_id: str, system_prompt: str, user_content: Union[str, List[Dict]]):
+    """
+    Processes a single user message and returns the AI response,
+    persisting memory via the thread_id.
+    """
+    config = {
+        "configurable": {
+            "thread_id": thread_id,
+            "system_prompt": system_prompt
+        }
+    }
+
+    # Prepare the input for this specific turn
+    input_data = {"messages": [HumanMessage(content=user_content)]}
+
+    ai_response = ""
+
+    # Stream the values to get the final AI message
+    for event in app.stream(input_data, config, stream_mode="values"):
+        if "messages" in event:
+            last_msg = event["messages"][-1]
+            if last_msg.type == "ai":
+                ai_response = last_msg.content
+
+    return ai_response
+
+# 使用范例
+# if __name__ == "__main__":
+#     tid = "py_expert_001"
+#     sys_p = "你是个善解人意的机器人。"
+
+#     # Call 1: Establish context
+#     resp1 = chat(tid, sys_p, "你好，我叫小明。")
+#     print(f"Bot: {resp1}")
+
+#     # Call 2: Test memory (The model should remember the name '小明')
+#     resp2 = chat(tid, sys_p, "我今天很开心")
+#     print(f"Bot: {resp2}")