MCP-сервер для управления памятью агента: memory_client.py

solutions/69f8e929da860fb4533faa2a_MCP-сервер_для_управления_памятью_агента/memory_client.py

@@ -1,96 +1,104 @@
 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) -> dict:
+def save(key: str, value: str) -> None:
     """Сохраняет значение по ключу."""
-    resp = httpx.post(f"{BASE_URL}/save", json={"key": key, "value": value})
-    return resp.json()
+    response = httpx.post(f"{BASE_URL}/save", json={"key": key, "value": value})
+    response.raise_for_status()
+    console.print(f"[green]✅[/green] Сохранено: {key} → {value}")
 
 
-def get(key: str) -> dict:
+def get(key: str) -> None:
     """Получает значение по ключу."""
-    resp = httpx.get(f"{BASE_URL}/get", params={"key": key})
-    return resp.json()
+    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']}")
 
 
-def delete(key: str) -> dict:
-    """Удаляет запись по ключу."""
-    resp = httpx.delete(f"{BASE_URL}/delete", params={"key": key})
-    return resp.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 list_keys() -> dict:
-    """Возвращает список всех ключей."""
-    resp = httpx.get(f"{BASE_URL}/list_keys")
-    return resp.json()
+def list_keys() -> None:
+    """Выводит список всех ключей."""
+    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)
 
 
-def save_with_namespace(namespace: str, key: str, value: str) -> dict:
+def save_with_namespace(namespace: str, key: str, value: str) -> None:
     """Сохраняет значение в указанном пространстве имён."""
-    resp = httpx.post(
+    response = httpx.post(
         f"{BASE_URL}/save_with_namespace",
         json={"namespace": namespace, "key": key, "value": value},
     )
-    return resp.json()
+    response.raise_for_status()
+    console.print(f"[green]✅[/green] Сохранено: {namespace}:{key} → {value}")
 
 
-def get_by_namespace(namespace: str, key: str) -> dict:
+def get_by_namespace(namespace: str, key: str) -> None:
     """Получает значение из указанного пространства имён."""
-    resp = httpx.get(
+    response = httpx.get(
         f"{BASE_URL}/get_by_namespace",
         params={"namespace": namespace, "key": key},
     )
-    return resp.json()
+    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']}")
 
 
 def main() -> None:
-    print("=== Тестирование MCP‑сервера ===")
-    # 1. Сохраняем пару ключ‑значение
-    print("\nСохранение 'foo':'bar'")
-    print(json.dumps(save("foo", "bar"), indent=2))
+    """Тестовый сценарий работы с сервером."""
+    # Очистка тестовых данных (если они уже есть)
+    delete("test_key")
+    delete("user1:session")
 
-    # 2. Получаем значение по ключу
-    print("\nПолучаем 'foo'")
-    print(json.dumps(get("foo"), indent=2))
+    console.print("[bold underline]=== Тестирование сервера памяти ===\n[/bold underline]")
 
-    # 3. Список всех ключей
-    print("\nСписок ключей")
-    print(json.dumps(list_keys(), indent=2))
+    # 1. Сохраняем и получаем обычный ключ
+    save("test_key", "Hello, World!")
+    get("test_key")
 
-    # 4. Удаляем ключ
-    print("\nУдаляем 'foo'")
-    print(json.dumps(delete("foo"), indent=2))
+    # 2. Удаляем ключ и проверяем отсутствие
+    delete("test_key")
+    get("test_key")
 
-    # 5. Проверяем, что ключ удалён
-    print("\nПолучаем 'foo' после удаления")
-    print(json.dumps(get("foo"), indent=2))
+    # 3. Работа с пространством имён
+    save_with_namespace("user1", "session", "abc123")
+    get_by_namespace("user1", "session")
 
-    # 6. Работа с пространством имён
-    ns = "experiment"
-    print(f"\nСохраняем в namespace '{ns}' key='x', value='123'")
-    print(json.dumps(save_with_namespace(ns, "x", "123"), indent=2))
+    # 4. Список всех ключей (должен включать только user1:session)
+    list_keys()
 
-    print("\nПолучаем из namespace")
-    print(json.dumps(get_by_namespace(ns, "x"), indent=2))
-
-    # 7. Список ключей (должен быть пустой после удаления)
-    print("\nСписок ключей после всех операций")
-    print(json.dumps(list_keys(), indent=2))
+    console.print("\n[bold green]✅ Тесты завершены успешно![/bold green]")
 
 
 if __name__ == "__main__":
-    try:
     main()
-    except httpx.RequestError as exc:
-        print(f"Ошибка соединения с сервером: {exc}")
-        sys.exit(1)
-    except Exception as e:
-        print(f"Неожиданная ошибка: {e}")
-        sys.exit(1)