107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
import json
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List
|
||
|
||
# ---------- FileCheckpointSaver ----------
|
||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||
|
||
class FileCheckpointSaver(BaseCheckpointSaver):
|
||
"""Сохраняет чекпоинт в JSON-файл."""
|
||
|
||
def __init__(self, filepath: str = "./checkpoint.json"):
|
||
self.filepath = Path(filepath)
|
||
self.filepath.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
def get(self, config: Dict[str, Any]) -> Dict[str, Any] | None:
|
||
if self.filepath.exists():
|
||
with open(self.filepath, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
return None
|
||
|
||
def put(self, config: Dict[str, Any], checkpoint: Dict[str, Any]) -> Dict[str, Any]:
|
||
with open(self.filepath, "w", encoding="utf-8") as f:
|
||
json.dump(checkpoint, f, indent=2, ensure_ascii=False)
|
||
return config
|
||
|
||
def list(self, config: Dict[str, Any]) -> List[str]:
|
||
if self.filepath.exists():
|
||
return [self.filepath.name]
|
||
return []
|
||
|
||
# ---------- ConversationMemory ----------
|
||
class ConversationMemory:
|
||
"""Управляет историей разговора с файловой persistence."""
|
||
|
||
def __init__(self, filepath: str = "./memory.json"):
|
||
self.filepath = Path(filepath)
|
||
self.history: List[Dict[str, str]] = self._load()
|
||
|
||
def _load(self) -> List[Dict[str, str]]:
|
||
if self.filepath.exists():
|
||
with open(self.filepath, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
return []
|
||
|
||
def add(self, role: str, content: str) -> None:
|
||
self.history.append({"role": role, "content": content})
|
||
self._save()
|
||
|
||
def get_history(self, limit: int = 10) -> List[Dict[str, str]]:
|
||
return self.history[-limit:]
|
||
|
||
def _save(self) -> None:
|
||
self.filepath.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(self.filepath, "w", encoding="utf-8") as f:
|
||
json.dump(self.history, f, indent=2, ensure_ascii=False)
|
||
|
||
def clear(self) -> None:
|
||
self.history = []
|
||
self._save()
|
||
|
||
# ---------- LangGraph Agent ----------
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.graph import add_messages
|
||
from typing import Annotated, TypedDict
|
||
|
||
class AgentState(TypedDict):
|
||
messages: Annotated[list, add_messages]
|
||
memory_summary: str
|
||
|
||
# Simple LLM placeholder – replace with real LLM
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
llm = ChatOpenAI(temperature=0.7)
|
||
|
||
|
||
def agent_node(state: AgentState, llm) -> AgentState:
|
||
response = llm.invoke(state["messages"])
|
||
return {"messages": [response]}
|
||
|
||
# Build graph
|
||
builder = StateGraph(AgentState)
|
||
builder.add_node("agent", agent_node)
|
||
builder.set_entry_point("agent")
|
||
builder.add_edge(START, "agent")
|
||
builder.add_edge("agent", END)
|
||
agent = builder.compile()
|
||
|
||
# ---------- CLI ----------
|
||
|
||
def chat_loop(agent, memory: ConversationMemory):
|
||
thread_id = "default"
|
||
while True:
|
||
user_input = input("\nВы: ")
|
||
if user_input.lower() in ["exit", "quit"]:
|
||
break
|
||
memory.add("user", user_input)
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
result = agent.invoke({"messages": [{"role": "human", "content": user_input}]}, config=config)
|
||
assistant_message = result["messages"][-1].content
|
||
memory.add("assistant", assistant_message)
|
||
print(f"\nАгент: {assistant_message}")
|
||
|
||
if __name__ == "__main__":
|
||
checkpointer = FileCheckpointSaver("./checkpoint.json")
|
||
memory = ConversationMemory("./memory.json")
|
||
chat_loop(agent, memory)
|