Server Introduction

import asyncio
import os
import sys
from datetime import datetime, timezone
from typing import Any

# Ensure we can import memmachine
sys.path.append(os.getcwd())

from memmachine.main.memmachine import MemMachine, MemoryType
from memmachine.common.episode_store import EpisodeEntry
from memmachine.server.api_v2.mcp import initialize_resource

class SessionDataWrapper:
    """Wrapper to adapt SessionInfo to SessionData protocol."""
    def __init__(self, session_info: Any):
        self._info = session_info

    @property
    def session_key(self) -> str:
        return self._info.episode_memory_conf.session_key

    @property
    def user_profile_id(self) -> str | None:
        return None

    @property
    def role_profile_id(self) -> str | None:
        return None

    @property
    def session_id(self) -> str | None:
        return self.session_key

async def main():
    config_file = "configuration.yml"
    if not os.path.exists(config_file):
        print(f"Error: Configuration file '{config_file}' not found in the current directory.")
        print("Please ensure you are running this script from the project root and that 'configuration.yml' exists.")
        sys.exit(1)

    print(f"Initializing MemMachine using {config_file}...")
    
    try:
        # Initialize MemMachine using the helper from mcp module which handles config loading
        mem_machine = await initialize_resource(config_file)
        
        # Create or get session
        session_key = "hello_session"
        print(f"Checking for session: {session_key}")
        
        session_info = await mem_machine.get_session(session_key=session_key)
        if session_info:
            print(f"Session {session_key} already exists, using it.")
        else:
            print(f"Creating session: {session_key}")
            # Note: create_session might fail if DBs are not reachable
            session_info = await mem_machine.create_session(session_key=session_key)
        
        # Wrap session info to satisfy SessionData protocol
        session = SessionDataWrapper(session_info)
        
        # Add a memory
        print("Adding memory...")
        episode = EpisodeEntry(
            content="Hello World! I am learning MemMachine.",
            producer_id="user1",
            producer_role="user",
            created_at=datetime.now(timezone.utc)
        )
        await mem_machine.add_episodes(
            session_data=session,
            episode_entries=[episode],
            target_memories=[MemoryType.Episodic]
        )
        
        # Search memory
        print("Searching memory...")
        results = await mem_machine.query_search(
            session_data=session,
            query="What am I learning?",
            limit=1
        )
        
        found = False
        if results.episodic_memory:
            # Check short-term memory
            if results.episodic_memory.short_term_memory and results.episodic_memory.short_term_memory.episodes:
                print(f"Found in Short Term Memory: {results.episodic_memory.short_term_memory.episodes[0].content}")
                found = True
            
            # Check long-term memory
            if results.episodic_memory.long_term_memory and results.episodic_memory.long_term_memory.episodes:
                print(f"Found in Long Term Memory: {results.episodic_memory.long_term_memory.episodes[0].content}")
                found = True
                
        if not found:
            print("No memory found (might need time to index or mock setup)")

        # Cleanup session
        print("Cleaning up session...")
        await mem_machine.delete_session(session)
        await mem_machine.stop()

    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"An error occurred: {e}")
        print("Ensure Neo4j is running and API keys are set in configuration.yml.")

if __name__ == "__main__":
    asyncio.run(main())