Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add authentication middleware to REST API #898

Open
cpacker opened this issue Jan 22, 2024 · 1 comment
Open

Add authentication middleware to REST API #898

cpacker opened this issue Jan 22, 2024 · 1 comment
Assignees
Labels
API Related to MemGPT API ChatUI Chat Interface

Comments

@cpacker
Copy link
Owner

cpacker commented Jan 22, 2024

Add optional authentication middleware to REST API, allowing clients to authenticate with bearer tokens

@cpacker cpacker self-assigned this Jan 22, 2024
@cpacker cpacker added API Related to MemGPT API ChatUI Chat Interface labels Jan 22, 2024
@arduenify
Copy link
Contributor

@cpacker What do you think of this implementation? Also, we can modify token storage from in-memory to persistent if desired.

Changes

server.py (SyncServer)

def authenticate_user(self) -> str:
    """
    Generates a secure random bearer token, creates a new user if necessary,
    and stores the token associated with the user.
    """
    user_id = uuid.uuid4() 
    token = secrets.token_urlsafe()
    user = User(id=user_id)
        
    try:
        self.ms.create_user(user)
    except ValueError:
        # user already exists
        pass
    self.active_tokens[token] = user_id
    return token

def verify_token(self, token: str) -> Optional[uuid.UUID]:
    return self.active_tokens.get(token, None)

auth/index.py

security = HTTPBearer()

class AuthResponse(BaseModel):
    token: str = Field(..., description="Bearer token for the authenticated user")


def setup_auth_router(server: SyncServer, interface: QueuingInterface):
    @router.get("/auth", tags=["auth"], response_model=AuthResponse)
    def authenticate_user():
        """
        Authenticates the user and sends response with User related data.

        Now returns a bearer token for the authenticated user.
        """
        interface.clear()
        try:
            token = server.authenticate_user()
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        return AuthResponse(token=token)
    
    def get_current_user(credentials: HTTPAuthorizationCredentials = Security(security)):
        token = credentials.credentials
        user_id = server.verify_token(token)
        if not user_id:
            raise HTTPException(status_code=403, detail="Invalid authentication credentials")
        user = server.ms.get_user(user_id)
        if not user:
            raise HTTPException(status_code=404, detail="User not found")
        return user

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
API Related to MemGPT API ChatUI Chat Interface
Projects
Development

No branches or pull requests

2 participants