MCP-сервер для управления памятью агента: memory_client.py
This commit is contained in:
+55
-72
@@ -1,104 +1,87 @@
|
||||
import json
|
||||
<|channel|>final code<|message|>import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
|
||||
def save(key: str, value: str) -> None:
|
||||
BASE_URL = "http://localhost:8000"
|
||||
|
||||
|
||||
def save(key: str, value: str) -> dict:
|
||||
"""Сохраняет значение по ключу."""
|
||||
response = httpx.post(f"{BASE_URL}/save", json={"key": key, "value": value})
|
||||
response.raise_for_status()
|
||||
console.print(f"[green]✅[/green] Сохранено: {key} → {value}")
|
||||
payload = {"key": key, "value": value}
|
||||
response = httpx.post(f"{BASE_URL}/save", json=payload)
|
||||
return response.json()
|
||||
|
||||
|
||||
def get(key: str) -> None:
|
||||
def get(key: str) -> dict:
|
||||
"""Получает значение по ключу."""
|
||||
response = httpx.get(f"{BASE_URL}/get", params={"key": key})
|
||||
if response.status_code == 404:
|
||||
console.print(f"[red]❌[/red] Ключ не найден: {key}")
|
||||
return
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
console.print(f"[blue]📥[/blue] Получено: {data['value']}")
|
||||
params = {"key": key}
|
||||
response = httpx.get(f"{BASE_URL}/get", params=params)
|
||||
return response.json()
|
||||
|
||||
|
||||
def delete(key: str) -> None:
|
||||
"""Удаляет ключ."""
|
||||
response = httpx.delete(f"{BASE_URL}/delete", params={"key": key})
|
||||
if response.status_code == 404:
|
||||
console.print(f"[red]❌[/red] Ключ не найден: {key}")
|
||||
return
|
||||
response.raise_for_status()
|
||||
console.print(f"[yellow]⚠️[/yellow] Удалено: {key}")
|
||||
def delete(key: str) -> dict:
|
||||
"""Удаляет запись по ключу."""
|
||||
payload = {"key": key}
|
||||
response = httpx.delete(f"{BASE_URL}/delete", json=payload)
|
||||
return response.json()
|
||||
|
||||
|
||||
def list_keys() -> None:
|
||||
"""Выводит список всех ключей."""
|
||||
def list_keys() -> dict:
|
||||
"""Возвращает список всех ключей."""
|
||||
response = httpx.get(f"{BASE_URL}/list_keys")
|
||||
response.raise_for_status()
|
||||
keys = response.json().get("keys", [])
|
||||
table = Table(title="Список ключей")
|
||||
table.add_column("№", justify="right")
|
||||
table.add_column("Ключ")
|
||||
for idx, key in enumerate(keys, 1):
|
||||
table.add_row(str(idx), key)
|
||||
console.print(table)
|
||||
return response.json()
|
||||
|
||||
|
||||
def save_with_namespace(namespace: str, key: str, value: str) -> None:
|
||||
def save_with_namespace(namespace: str, key: str, value: str) -> dict:
|
||||
"""Сохраняет значение в указанном пространстве имён."""
|
||||
response = httpx.post(
|
||||
f"{BASE_URL}/save_with_namespace",
|
||||
json={"namespace": namespace, "key": key, "value": value},
|
||||
)
|
||||
response.raise_for_status()
|
||||
console.print(f"[green]✅[/green] Сохранено: {namespace}:{key} → {value}")
|
||||
payload = {"namespace": namespace, "key": key, "value": value}
|
||||
response = httpx.post(f"{BASE_URL}/save_with_namespace", json=payload)
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_by_namespace(namespace: str, key: str) -> None:
|
||||
def get_by_namespace(namespace: str, key: str) -> dict:
|
||||
"""Получает значение из указанного пространства имён."""
|
||||
response = httpx.get(
|
||||
f"{BASE_URL}/get_by_namespace",
|
||||
params={"namespace": namespace, "key": key},
|
||||
)
|
||||
if response.status_code == 404:
|
||||
console.print(f"[red]❌[/red] Ключ не найден: {namespace}:{key}")
|
||||
return
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
console.print(f"[blue]📥[/blue] Получено из {namespace}: {data['value']}")
|
||||
params = {"namespace": namespace, "key": key}
|
||||
response = httpx.get(f"{BASE_URL}/get_by_namespace", params=params)
|
||||
return response.json()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Тестовый сценарий работы с сервером."""
|
||||
# Очистка тестовых данных (если они уже есть)
|
||||
delete("test_key")
|
||||
delete("user1:session")
|
||||
"""Простой тестовый сценарий для проверки работы сервера."""
|
||||
print("=== Тестирование MCP-сервера ===")
|
||||
|
||||
console.print("[bold underline]=== Тестирование сервера памяти ===\n[/bold underline]")
|
||||
# 1. Сохраняем простое значение
|
||||
print("\nСохраняем ключ 'foo' со значением 'bar'")
|
||||
print(json.dumps(save("foo", "bar"), indent=2))
|
||||
|
||||
# 1. Сохраняем и получаем обычный ключ
|
||||
save("test_key", "Hello, World!")
|
||||
get("test_key")
|
||||
# 2. Получаем сохранённое значение
|
||||
print("\nПолучаем ключ 'foo'")
|
||||
print(json.dumps(get("foo"), indent=2))
|
||||
|
||||
# 2. Удаляем ключ и проверяем отсутствие
|
||||
delete("test_key")
|
||||
get("test_key")
|
||||
# 3. Сохраняем в пространстве имён
|
||||
ns = "experiment"
|
||||
print(f"\nСохраняем ключ '{ns}:baz' со значением 'qux' в namespace '{ns}'")
|
||||
print(json.dumps(save_with_namespace(ns, "baz", "qux"), indent=2))
|
||||
|
||||
# 3. Работа с пространством имён
|
||||
save_with_namespace("user1", "session", "abc123")
|
||||
get_by_namespace("user1", "session")
|
||||
# 4. Получаем из пространства имён
|
||||
print(f"\nПолучаем ключ '{ns}:baz' из namespace '{ns}'")
|
||||
print(json.dumps(get_by_namespace(ns, "baz"), indent=2))
|
||||
|
||||
# 4. Список всех ключей (должен включать только user1:session)
|
||||
list_keys()
|
||||
# 5. Список всех ключей
|
||||
print("\nСписок всех ключей:")
|
||||
print(json.dumps(list_keys(), indent=2))
|
||||
|
||||
console.print("\n[bold green]✅ Тесты завершены успешно![/bold green]")
|
||||
# 6. Удаляем ключ 'foo'
|
||||
print("\nУдаляем ключ 'foo'")
|
||||
print(json.dumps(delete("foo"), indent=2))
|
||||
|
||||
# 7. Проверяем, что ключ удалён
|
||||
print("\nПроверка удаления ключа 'foo':")
|
||||
print(json.dumps(get("foo"), indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user