diff --git a/main.py b/main.py index 41bc3cd..0651b7d 100644 --- a/main.py +++ b/main.py @@ -1,14 +1,15 @@ import os import asyncio -from typing import Any, Dict +from typing import Any -from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage -from langchain.tools import tool +from fastmcp import Client from deepagents import create_deep_agent from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend +from langchain_openai import ChatOpenAI +from langchain.tools import tool -# LLM via OpenRouter + +# LLM configuration (OpenRouter) llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -16,6 +17,7 @@ llm = ChatOpenAI( temperature=0.0, ) + backend = CompositeBackend( [ LocalShellBackend(workspace_dir="./workspace"), @@ -23,48 +25,58 @@ backend = CompositeBackend( ] ) -# Tool that forwards calls to the MCP memory server + @tool -def memory_action(action: str, params: Dict[str, Any]) -> str: - """ - Perform a memory operation via the MCP server. - - action: one of save_with_namespace, get_by_namespace, list_keys, delete, get, save - params: dictionary of parameters required by the chosen action. - Returns a JSON string with the server response. - """ - import json - import asyncio - from fastmcp import Client - +def memory_save(key: str, value: Any, namespace: str = "default") -> bool: + """Save a value in the remote memory server.""" async def _call(): client = Client("python memory_server.py") await client.connect() try: - result = await client.call_tool(action, params) - return json.dumps(result, ensure_ascii=False) + result = await client.call_tool( + "save_with_namespace", + {"key": key, "value": value, "namespace": namespace}, + ) + return result finally: await client.close() + return asyncio.run(_call()) - loop = asyncio.get_event_loop() - return loop.run_until_complete(_call()) + +@tool +def memory_get(namespace: str = "default") -> str: + """Retrieve all key-value pairs from a namespace as a formatted string.""" + async def _call(): + client = Client("python memory_server.py") + await client.connect() + try: + data = await client.call_tool("get_by_namespace", {"namespace": namespace}) + if not data: + return "No data." + lines = [f"{item['key']}: {item['value']}" for item in data] + return "\\n".join(lines) + finally: + await client.close() + return asyncio.run(_call()) agent = create_deep_agent( model=llm, - tools=[memory_action], + tools=[memory_save, memory_get], backend=backend, - system_prompt="You are an assistant that can store and retrieve data using a remote memory server.", + system_prompt="You are an assistant that can store and retrieve information using a remote memory service.", ) + async def demo(): - # Example: store a name and then retrieve it - query = ( - "Save the user name 'Алексей' in the default namespace using the memory_action tool. " - "Then read back all entries from the default namespace and return them." + # Store a fact + await agent.ainvoke( + {"messages": [{"role": "user", "content": "Запомни, что мой любимый цвет - синий."}]}, + {"configurable": {"thread_id": "demo-1"}}, ) + # Retrieve stored facts result = await agent.ainvoke( - {"messages": [HumanMessage(content=query)]}, + {"messages": [{"role": "user", "content": "Что я просил запомнить?"}]}, {"configurable": {"thread_id": "demo-1"}}, ) print(result["messages"][-1].content)