diff --git a/agent_core.py b/agent_core.py index 40548ce..1d885a7 100644 --- a/agent_core.py +++ b/agent_core.py @@ -27,6 +27,7 @@ from llm_router import get_llm from memory import add_memory, search_memory, read_cold_storage, load_profile +from tenant import is_guest from skills.email_manager import (run_email_summary, draft_email_reply, update_memory, read_email_thread) from people import remember_person, get_person, list_people @@ -80,8 +81,15 @@ def generate_morning_news() -> str: return "News briefing generated successfully." -def build_tools(): - """The full tool array Aria can call. Add new skills here (see CONTRIBUTING.md).""" +def build_tools(guest=False): + """The tool array Aria can call. Add new skills here (see CONTRIBUTING.md). + + guest=True returns the restricted, account-free subset for friend/guest sessions: + isolated per-user memory + stateless general tools ONLY. Nothing that touches the + owner's accounts or shared data (no email, calendar, contacts, notes, reminders, + music, smart-home, browser, or system access).""" + if guest: + return [add_memory, search_memory, get_weather, web_search, fetch_webpage] return [ add_memory, search_memory, @@ -140,7 +148,7 @@ def _stable_prompt() -> str: standing rule or profile (rare), which is the right time to rebuild the cache. """ profile_data = load_profile() - standing_instructions = render_for_prompt() + standing_instructions = "" if is_guest() else render_for_prompt() return f"""You are Aria (Aria Responds Intelligently Always), a highly intelligent, proactive, and friendly personal AI assistant. Your goal is to help your user manage their life, emails, and news. @@ -277,8 +285,18 @@ def build_system_message() -> SystemMessage: # assistant interactions — the dominant cause of cache misses for personal use. # Reads stay ~0.1x; the write premium (2x vs 1.25x) pays off after one re-hit/hour. ttl = os.getenv("ARIA_CACHE_TTL", "1h") + stable = _stable_prompt() + if is_guest(): + # Guest sessions get a restricted toolset; tell the model so it sets expectations + # instead of pretending to have email/calendar. (Profile + standing rules are + # already empty for guests, so this banner is the only owner-vs-guest prompt delta.) + stable = ("GUEST MODE — you are a personal demo of Aria for a guest user. You have " + "ONLY: your memory of THIS guest (add_memory / search_memory), web search, " + "and weather. You do NOT have email, calendar, contacts, notes, reminders, " + "music, smart-home, or anyone's accounts — if asked for those, say you " + "can't in guest mode. Be warm, and remember what they tell you.\n\n") + stable return SystemMessage(content=[ - {"type": "text", "text": _stable_prompt(), + {"type": "text", "text": stable, "cache_control": {"type": "ephemeral", "ttl": ttl}}, {"type": "text", "text": f"The current date and time is: {now}. Use this for any date math " @@ -306,16 +324,17 @@ def thread_config(thread_id: str) -> dict: return {"configurable": {"thread_id": thread_id}} -def build_agent(checkpointer=None): +def build_agent(checkpointer=None, guest=False): """Construct the LangGraph ReAct agent. Raises if no LLM can be initialized. Pass a checkpointer (open_checkpointer()) to get durable conversations; omit it - for stateless one-shot use. + for stateless one-shot use. guest=True builds the restricted guest agent (account-free + toolset); invoke it only with a tenant context set (see tenant.set_current_user). """ llm = get_llm(temperature=0) return create_agent( llm, - build_tools(), + build_tools(guest=guest), middleware=[_fresh_system_prompt], checkpointer=checkpointer, ) diff --git a/memory.py b/memory.py index 430c2b7..e2c59f8 100644 --- a/memory.py +++ b/memory.py @@ -13,6 +13,8 @@ from langchain_google_genai import GoogleGenerativeAIEmbeddings from dotenv import load_dotenv +from tenant import is_guest, get_current_user, safe_id + load_dotenv() PROFILE_PATH = os.path.join(os.path.dirname(__file__), "profile.json") @@ -22,6 +24,9 @@ # --- Layer 1: Core Identity (Static Profile) --- def load_profile(): + # Guests have no static owner profile — their personalization lives in their own memory. + if is_guest(): + return {} if not os.path.exists(PROFILE_PATH): return {} with open(PROFILE_PATH, "r") as f: @@ -46,6 +51,7 @@ def update_profile(key, value): # --- Layer 2: Semantic Memory (ChromaDB) --- # Initialize ChromaDB correctly on first load +chroma_client = None try: chroma_client = chromadb.PersistentClient(path=DB_PATH) collection = chroma_client.get_or_create_collection(name="aria_memory") @@ -58,11 +64,30 @@ def update_profile(key, value): collection = None embeddings = None + +def _guest_collection(): + """The current guest's OWN ChromaDB collection — isolated per user (`mem_`).""" + if chroma_client is None: + return None + return chroma_client.get_or_create_collection(name=f"mem_{safe_id(get_current_user())}") + @tool def add_memory(fact: str) -> str: """Use this tool to add a new memory, fact, preference, or event about the user. Provide a clear, detailed sentence (e.g., 'Satvik hates newsletters', 'Satvik is traveling to NY in July'). """ + if is_guest(): + # Guests write straight into their own isolated collection (no shared scratchpad, + # no compaction job — those are owner-only). + col = _guest_collection() + if col is None or embeddings is None: + return "Memory isn't available right now." + try: + col.add(ids=[uuid.uuid4().hex], documents=[fact], + embeddings=[embeddings.embed_query(fact)]) + return f"Got it — I'll remember that: {fact}" + except Exception as e: + return f"Failed to add memory: {str(e)}" try: with open(SCRATCHPAD_PATH, "a") as f: f.write(fact + "\n") @@ -76,8 +101,22 @@ def search_memory(query: str, n_results: int = 3) -> str: """Use this tool to search the user's semantic memory for relevant facts or past events. Provide a search query (e.g., 'Does Satvik like coffee?', 'travel plans'). """ + if is_guest(): + # Query ONLY this guest's collection — never the owner's scratchpad/collection. + col = _guest_collection() + if col is None or embeddings is None: + return "No relevant memories found." + try: + res = col.query(query_embeddings=[embeddings.embed_query(query)], n_results=n_results) + docs = res['documents'][0] if res.get('documents') else [] + except Exception as e: + return f"[memory error: {str(e)}]" + if not docs: + return "No relevant memories found." + return "What I remember:\n" + "\n".join(f"- {d}" for d in docs) + output = [] - + # 1. Working Memory (Scratchpad) if os.path.exists(SCRATCHPAD_PATH): with open(SCRATCHPAD_PATH, "r") as f: diff --git a/tenant.py b/tenant.py new file mode 100644 index 0000000..bca9414 --- /dev/null +++ b/tenant.py @@ -0,0 +1,44 @@ +"""Per-request tenant context — the seam that lets one Aria serve isolated guest users. + +The OWNER (you) runs with NO tenant set: everything behaves exactly as before — your +profile, your memory collection, your scratchpad, your full toolset. + +A GUEST (a friend trying it out) is set as the current user for the duration of a request. +While set, memory routes to that guest's OWN ChromaDB collection (`mem_`), the static +profile is empty, and the agent is built with a restricted, account-free toolset. Friends +can never see your data or each other's. + +Usage (the caller owns set/reset, always in a try/finally so context can't leak): + token = set_current_user("alice") + try: + agent.invoke(...) + finally: + reset_current_user(token) +""" +import re +import contextvars + +_current_user: contextvars.ContextVar = contextvars.ContextVar("current_user", default=None) + + +def set_current_user(user_id): + """Set the current guest for this context. Returns a token for reset_current_user().""" + return _current_user.set(user_id) + + +def reset_current_user(token): + _current_user.reset(token) + + +def get_current_user(): + return _current_user.get() + + +def is_guest() -> bool: + """True when a guest user is set (i.e. NOT the owner / default context).""" + return _current_user.get() is not None + + +def safe_id(user_id: str) -> str: + """A ChromaDB-safe collection suffix for a user id (alphanumerics + underscore).""" + return re.sub(r"[^a-zA-Z0-9]", "_", str(user_id))[:48] or "anon" diff --git a/tests/test_tenant.py b/tests/test_tenant.py new file mode 100644 index 0000000..157927f --- /dev/null +++ b/tests/test_tenant.py @@ -0,0 +1,98 @@ +"""Phase-1 multi-tenancy: tenant context, guest toolset, and per-user memory isolation. + +The actual ChromaDB add/search needs embeddings (network), so those paths are mocked — +we assert the routing (which collection) rather than real vector recall. Offline. +Run: python3 -m unittest discover tests +""" +import unittest +from unittest.mock import patch, MagicMock, mock_open + +import tenant +import memory +import agent_core + + +class _GuestCtx: + """with _GuestCtx('alice'): … — sets/clears the tenant context safely.""" + def __init__(self, uid): + self.uid = uid + def __enter__(self): + self.tok = tenant.set_current_user(self.uid) + def __exit__(self, *a): + tenant.reset_current_user(self.tok) + + +class TestTenantContext(unittest.TestCase): + def test_default_is_owner(self): + self.assertIsNone(tenant.get_current_user()) + self.assertFalse(tenant.is_guest()) + + def test_set_and_reset(self): + with _GuestCtx("alice"): + self.assertEqual(tenant.get_current_user(), "alice") + self.assertTrue(tenant.is_guest()) + self.assertFalse(tenant.is_guest()) # restored after the block + + def test_safe_id(self): + self.assertEqual(tenant.safe_id("alice@x.com"), "alice_x_com") + self.assertEqual(tenant.safe_id(""), "anon") + + +class TestGuestToolset(unittest.TestCase): + def test_guest_excludes_account_tools(self): + names = {t.name for t in agent_core.build_tools(guest=True)} + for keep in ("add_memory", "search_memory", "web_search", "fetch_webpage", "get_weather"): + self.assertIn(keep, names) + for forbidden in ("read_email_thread", "read_and_summarize_emails", "draft_email_reply", + "get_calendar_events", "create_calendar_event", "list_people", + "play_music", "control_light", "check_packages", "get_system_status", + "browse_and_report", "add_commitment", "create_note", "read_cold_storage"): + self.assertNotIn(forbidden, names) + + def test_owner_toolset_is_full(self): + names = {t.name for t in agent_core.build_tools(guest=False)} + self.assertIn("read_email_thread", names) + self.assertGreater(len(names), 30) + + +class TestGuestMemoryIsolation(unittest.TestCase): + def _mocks(self, query_docs=None): + client, col, emb = MagicMock(), MagicMock(), MagicMock() + client.get_or_create_collection.return_value = col + emb.embed_query.return_value = [0.1, 0.2] + col.query.return_value = {'documents': [query_docs or []]} + return client, col, emb + + def test_guest_add_routes_to_own_collection(self): + client, col, emb = self._mocks() + with _GuestCtx("bob"), patch.object(memory, 'chroma_client', client), \ + patch.object(memory, 'embeddings', emb): + out = memory.add_memory.invoke({'fact': 'bob likes tea'}) + client.get_or_create_collection.assert_called_once_with(name='mem_bob') + col.add.assert_called_once() + self.assertIn('remember', out.lower()) + + def test_owner_add_uses_scratchpad_not_chroma(self): + client, _, emb = self._mocks() + m = mock_open() + with patch('builtins.open', m), patch.object(memory, 'chroma_client', client): + out = memory.add_memory.invoke({'fact': 'owner fact'}) + client.get_or_create_collection.assert_not_called() # owner never hits guest path + m.assert_called() # wrote the scratchpad + self.assertIn('working memory', out.lower()) + + def test_guest_search_queries_only_own_collection(self): + client, col, emb = self._mocks(query_docs=['bob likes tea']) + with _GuestCtx("bob"), patch.object(memory, 'chroma_client', client), \ + patch.object(memory, 'embeddings', emb): + out = memory.search_memory.invoke({'query': 'drinks'}) + client.get_or_create_collection.assert_called_once_with(name='mem_bob') + self.assertIn('bob likes tea', out) + + def test_guest_profile_is_empty(self): + with _GuestCtx("bob"): + self.assertEqual(memory.load_profile(), {}) + + +if __name__ == '__main__': + unittest.main()