import os
import re
import json
from statebase import StateBase
from google import genai
from google.genai import types
# 1. Setup Clients
sb = StateBase(api_key=os.getenv("STATEBASE_KEY"))
gemini = genai.Client(api_key=os.getenv("GEMINI_KEY"))
# 2. Initialize Persistent Session
session = sb.sessions.create(
agent_id="senior-support-v1",
initial_state={"status": "ready", "user_name": "Guest"}
)
def process_interaction(user_input):
# 3. Get Context (State + Memories + History)
context = sb.sessions.get_context(session_id=session.id, query=user_input)
# 4. Prompt with State-Awareness
prompt = f"""
Current State: {context['state']}
Relevant Memories: {context['memories']}
If you learn something new, use <memory>...</memory>.
To update state, use <state_update>{{"key": "value"}}</state_update>.
"""
response = gemini.models.generate_content(
model="gemini-2.0-flash",
config=types.GenerateContentConfig(system_instruction=prompt),
contents=user_input
)
raw_output = response.text
# 5. Extract and Persist (The StateBase Magic)
# Memory extraction
mem_match = re.search(r"<memory>(.*?)</memory>", raw_output)
if mem_match:
sb.memory.add(content=mem_match.group(1), session_id=session.id)
# State extraction
state_match = re.search(r"<state_update>(.*?)</state_update>", raw_output)
if state_match:
new_state = json.loads(state_match.group(1))
sb.sessions.update_state(session_id=session.id, state=new_state)
# 6. Log Turn for Dashboard Trace
sb.sessions.add_turn(session_id=session.id, input=user_input, output=raw_output)
return re.sub(r"<(memory|state_update)>.*?</\1>", "", raw_output).strip()
# Run the loop
if __name__ == "__main__":
print(process_interaction("My name is Alex and I love python."))