Skip to content

Commit

Permalink
up
Browse files Browse the repository at this point in the history
  • Loading branch information
emrgnt-cmplxty committed Aug 23, 2024
1 parent f0e8369 commit 29f30ec
Show file tree
Hide file tree
Showing 16 changed files with 654 additions and 213 deletions.
95 changes: 44 additions & 51 deletions py/cli/commands/management.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,23 @@ def app_settings(client):

@cli.command()
@click.option("--user-ids", multiple=True, help="User IDs to overview")
@click.option(
"--offset",
default=None,
help="The offset to start from. Defaults to 0.",
)
@click.option(
"--limit",
default=None,
help="The maximum number of nodes to return. Defaults to 100.",
)
@click.pass_obj
def users_overview(client, user_ids):
def users_overview(client, user_ids, offset, limit):
"""Get an overview of users."""
user_ids = list(user_ids) if user_ids else None

with timer():
response = client.users_overview(user_ids)
response = client.users_overview(user_ids, offset, limit)

for user in response:
click.echo(user)
Expand Down Expand Up @@ -73,25 +83,45 @@ def delete(client, filter):

@cli.command()
@click.option("--document-ids", multiple=True, help="Document IDs to overview")
@click.option(
"--offset",
default=None,
help="The offset to start from. Defaults to 0.",
)
@click.option(
"--limit",
default=None,
help="The maximum number of nodes to return. Defaults to 100.",
)
@click.pass_obj
def documents_overview(client, document_ids):
def documents_overview(client, document_ids, offset, limit):
"""Get an overview of documents."""
document_ids = list(document_ids) if document_ids else None

with timer():
response = client.documents_overview(document_ids)
response = client.documents_overview(document_ids, offset, limit)

for document in response["results"]:
click.echo(document)


@cli.command()
@click.option("--document-id", help="Document ID to retrieve chunks for")
@click.option(
"--offset",
default=None,
help="The offset to start from. Defaults to 0.",
)
@click.option(
"--limit",
default=None,
help="The maximum number of nodes to return. Defaults to 100.",
)
@click.pass_obj
def document_chunks(client, document_id):
def document_chunks(client, document_id, offset, limit):
"""Get chunks of a specific document."""
with timer():
response = client.document_chunks(document_id)
response = client.document_chunks(document_id, offset, limit)

chunks = response.get("results", [])
click.echo(f"\nNumber of chunks: {len(chunks)}")
Expand All @@ -103,57 +133,20 @@ def document_chunks(client, document_id):


@cli.command()
@click.option(
"--offset",
default=None,
help="The offset to start from. Defaults to 0.",
)
@click.option(
"--limit",
default=None,
help="The maximum number of nodes to return. Defaults to 100.",
)
@click.pass_obj
def inspect_knowledge_graph(client, limit):
def inspect_knowledge_graph(client, offset, limit):
"""Inspect the knowledge graph."""
with timer():
response = client.inspect_knowledge_graph(limit)

click.echo(response)


## TODO: Implement groups_overview


## TODO: Implement create_group


## TODO: Implement get_group


## TODO: Implement update_group


## TODO: Implement delete_group


## TODO: Implement list_groups


## TODO: Implement add_user_to_group


## TODO: Implement remove_user_from_group


## TODO: Implement get_users_in_group


## TODO: Implement get_groups_for_user


## TODO: Implement assign_document_to_group


## TODO: Implement remove_document_from_group


## TODO: Implement document_groups

response = client.inspect_knowledge_graph(offset, limit)

## TODO: Implement get_documents_in_group
click.echo(response)
2 changes: 1 addition & 1 deletion py/core/examples/scripts/advanced_kg_cookbook.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ def main(

print("Inspecting Knowledge Graph")
print(
client.inspect_knowledge_graph(1000, print_descriptions=True)[
client.inspect_knowledge_graph(0, 1000, print_descriptions=True)[
"results"
]
)
Expand Down
164 changes: 164 additions & 0 deletions py/core/examples/scripts/run_group_workflow.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
import os
from r2r import R2RClient


if __name__ == "__main__":
# Initialize the R2R client
client = R2RClient("http://localhost:8000") # Replace with your R2R deployment URL

# Admin login
print("Logging in as admin...")
login_result = client.login("[email protected]", "change_me_immediately")
print("Admin login result:", login_result)

# Create two groups
print("\nCreating two groups...")
group1_result = client.create_group("TestGroup1", "A test group for document access")
group2_result = client.create_group("TestGroup2", "Another test group")
print("Group1 creation result:", group1_result)
print("Group2 creation result:", group2_result)
group1_id = group1_result['results']['group_id']
group2_id = group2_result['results']['group_id']

# Get groups overview
print("\nGetting groups overview...")
groups_overview = client.groups_overview()
print("Groups overview:", groups_overview)

# Get specific group
print("\nGetting specific group...")
group1_details = client.get_group(group1_id)
print("Group1 details:", group1_details)

# List all groups
print("\nListing all groups...")
groups_list = client.list_groups()
print("Groups list:", groups_list)

# Update a group
print("\nUpdating Group1...")
update_result = client.update_group(group1_id, name="UpdatedTestGroup1", description="Updated description")
print("Group update result:", update_result)

# Ingest two documents
print("\nIngesting two documents...")
script_path = os.path.dirname(__file__)
sample_file1 = os.path.join(script_path, "core", "examples", "data", "aristotle_v2.txt")
sample_file2 = os.path.join(script_path, "core", "examples", "data", "aristotle.txt")
ingestion_result1 = client.ingest_files([sample_file1])
ingestion_result2 = client.ingest_files([sample_file2])
print("Document1 ingestion result:", ingestion_result1)
print("Document2 ingestion result:", ingestion_result2)
document1_id = ingestion_result1['results']['processed_documents'][0]['id']
document2_id = ingestion_result2['results']['processed_documents'][0]['id']

# Assign documents to groups
print("\nAssigning documents to groups...")
assign_result1 = client.assign_document_to_group(document1_id, group1_id)
assign_result2 = client.assign_document_to_group(document2_id, group2_id)
print("Document1 assignment result:", assign_result1)
print("Document2 assignment result:", assign_result2)

# document1_id = "c3291abf-8a4e-5d9d-80fd-232ef6fd8526"
# Get document groups
print("\nGetting groups for Document1...")
doc1_groups = client.document_groups(document1_id)
print("Document1 groups:", doc1_groups)

# Create three test users
print("\nCreating three test users...")
user1_result = client.register("[email protected]", "password123")
user2_result = client.register("[email protected]", "password123")
user3_result = client.register("[email protected]", "password123")
print("User1 creation result:", user1_result)
print("User2 creation result:", user2_result)
print("User3 creation result:", user3_result)

# Add users to groups
print("\nAdding users to groups...")
add_user1_result = client.add_user_to_group(user1_result['results']['id'], group1_id)
add_user2_result = client.add_user_to_group(user2_result['results']['id'], group2_id)
add_user3_result1 = client.add_user_to_group(user3_result['results']['id'], group1_id)
add_user3_result2 = client.add_user_to_group(user3_result['results']['id'], group2_id)
print("Add user1 to group1 result:", add_user1_result)
print("Add user2 to group2 result:", add_user2_result)
print("Add user3 to group1 result:", add_user3_result1)
print("Add user3 to group2 result:", add_user3_result2)

# Get users in a group
print("\nGetting users in Group1...")
users_in_group1 = client.user_groups(group1_id)
print("Users in Group1:", users_in_group1)

# Get groups for a user
print("\nGetting groups for User3...")
user3_groups = client.user_groups(user3_result['results']['id'])
print("User3 groups:", user3_groups)

# Get documents in a group
print("\nGetting documents in Group1...")
docs_in_group1 = client.documents_in_group(group1_id)
print("Documents in Group1:", docs_in_group1)

# Remove user from group
print("\nRemoving User3 from Group1...")
remove_user_result = client.remove_user_from_group(user3_result['results']['id'], group1_id)
print("Remove user result:", remove_user_result)

# Remove document from group
print("\nRemoving Document1 from Group1...")
remove_doc_result = client.remove_document_from_group(document1_id, group1_id)
print("Remove document result:", remove_doc_result)

# Logout admin
print("\nLogging out admin...")
client.logout()

# Login as user1
print("\nLogging in as user1...")
client.login("[email protected]", "password123")

# Search for documents (should see document1 but not document2)
print("\nUser1 searching for documents...")
search_result_user1 = client.search("philosophy", {"selected_group_ids": [group1_id]})
print("User1 search result:", search_result_user1)

# Logout user1
print("\nLogging out user1...")
client.logout()

# Login as user3
print("\nLogging in as user3...")
client.login("[email protected]", "password123")

# Search for documents (should see only document2 after removal from Group1)
print("\nUser3 searching for documents...")
try:
search_result_user3 = client.search("philosophy", {"selected_group_ids": [group1_id, group2_id]})
except Exception as e:
print("User3 search result error:", e)
search_result_user3 = client.search("philosophy", {"selected_group_ids": [group2_id]})

print("User3 search result:", search_result_user3)

# Logout user3
print("\nLogging out user3...")
client.logout()

# Clean up
print("\nCleaning up...")
# Login as admin again
client.login("[email protected]", "change_me_immediately")

# Delete the groups
print("Deleting the groups...")
client.delete_group(group1_id)
client.delete_group(group2_id)

# Logout admin
print("\nLogging out admin...")
client.logout()

print("\nWorkflow completed.")


18 changes: 12 additions & 6 deletions py/core/main/api/routes/auth/base.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
import uuid
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional

from core.base.api.models.auth.responses import (
GenericMessageResponse,
WrappedGenericMessageResponse,
WrappedTokenResponse,
WrappedUserResponse,
)
from fastapi import Body, Depends
from fastapi import Path, Body, Depends
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from pydantic import EmailStr

Expand Down Expand Up @@ -200,14 +200,18 @@ async def reset_password_app(
return GenericMessageResponse(message=result["message"])

@self.router.delete(
"/user", response_model=WrappedGenericMessageResponse
"/user/{user_id}", response_model=WrappedGenericMessageResponse
)
@self.base_endpoint
async def delete_user_app(
user_id: str = Body(..., description="ID of the user to delete"),
password: str | None = Body(
user_id: str = Path(..., description="ID of the user to delete"),
password: Optional[str] = Body(
None, description="User's current password"
),
delete_vector_data: Optional[bool] = Body(
False,
description="Whether to delete the user's vector data",
),
auth_user=Depends(self.engine.providers.auth.auth_wrapper),
):
"""
Expand All @@ -218,6 +222,8 @@ async def delete_user_app(
"""
if auth_user.id != user_id and not auth_user.is_superuser:
raise Exception("User ID does not match authenticated user")
if not auth_user.is_superuser and not password:
raise Exception("Password is required for non-superusers")
user_uuid = uuid.UUID(user_id)
result = await self.engine.adelete_user(user_uuid, password)
result = await self.engine.adelete_user(user_uuid, password, delete_vector_data)
return GenericMessageResponse(message=result["message"])
Loading

0 comments on commit 29f30ec

Please sign in to comment.