import os
from openai import OpenAI
from atlas_mem import AtlasMem
openai_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
brain = AtlasMem(
api_key=os.environ["ATLAS_API_KEY"],
base_url="https://api.bsyncs.com",
user_id="user-123",
session_id="session-abc",
)
def chat(user_message: str) -> str:
# 1. Retrieve relevant memory
memory = brain.search(user_message, k=5)
memory_context = memory.format()
# 2. Build system prompt with memory
system_prompt = f"""You are a helpful assistant with long-term memory.
MEMORY CONTEXT (use this to personalise your response):
{memory_context}
Always reference specific facts from memory when relevant."""
# 3. Call the LLM
response = openai_client.chat.completions.create(
model="gpt-4o",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message},
]
)
assistant_reply = response.choices[0].message.content
# 4. Store both turns in memory
brain.add(user_message, source="user")
brain.add(assistant_reply, source="assistant")
return assistant_reply
# Example usage
print(chat("What database does Acme Corp use?"))