File size: 3,040 Bytes
316b3f1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6922a90
 
 
 
316b3f1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
#!/usr/bin/env python3
"""Local smoke test for the Smriti AI Hugging Face custom handler."""

from __future__ import annotations

import json
import os
from pathlib import Path

from handler import EndpointHandler


def pretty(title: str, payload: dict) -> None:
    print(f"\n=== {title} ===")
    print(json.dumps(payload, indent=2, ensure_ascii=False))


def main() -> int:
    # Default to memory-only smoke mode so local validation never downloads or
    # reports against a tiny/test model. Set BASE_MODEL_ID=google/gemma-4-E2B-it
    # or HF_ENDPOINT_URL explicitly to test real model generation.
    os.environ.setdefault("BASE_MODEL_ID", "")
    os.environ.setdefault("SMRITI_MEMORY_BACKEND", "json")
    os.environ.setdefault("SMRITI_MEMORY_PATH", "/tmp/smriti_hf_test.json")
    os.environ.setdefault("SMRITI_RETRIEVAL_MODE", "semantic_graph_identity")
    retrieval_mode = os.environ["SMRITI_RETRIEVAL_MODE"]

    handler = EndpointHandler(path=str(Path(__file__).resolve().parent))

    pretty("health", handler({"inputs": {"operation": "health"}}))

    pretty(
        "fact injection",
        handler(
            {
                "inputs": {
                    "operation": "chat",
                    "user_id": "local-demo-user",
                    "message": "My name is Alex and I am a marine biologist based in Hawaii.",
                    "retrieval_mode": retrieval_mode,
                },
                "parameters": {"max_new_tokens": 64, "return_memories": True},
            }
        ),
    )

    pretty(
        "distractor",
        handler(
            {
                "inputs": {
                    "operation": "chat",
                    "user_id": "local-demo-user",
                    "message": "What is the capital of France?",
                    "retrieval_mode": retrieval_mode,
                },
                "parameters": {"max_new_tokens": 64, "return_memories": True},
            }
        ),
    )

    pretty(
        "recall",
        handler(
            {
                "inputs": {
                    "operation": "chat",
                    "user_id": "local-demo-user",
                    "message": "What do you remember about me?",
                    "retrieval_mode": retrieval_mode,
                },
                "parameters": {"max_new_tokens": 64, "return_memories": True},
            }
        ),
    )

    pretty(
        "delete memory",
        handler({"inputs": {"operation": "delete_memory", "user_id": "local-demo-user"}}),
    )

    pretty(
        "recall after delete",
        handler(
            {
                "inputs": {
                    "operation": "chat",
                    "user_id": "local-demo-user",
                    "message": "What do you remember about me?",
                    "retrieval_mode": retrieval_mode,
                },
                "parameters": {"max_new_tokens": 64, "return_memories": True},
            }
        ),
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())