Add org & role (#226)

* Add organizataion & role

* update the frontend

Signed-off-by: kerthcet <kerthcet@gmail.com>

---------

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

.env.example

@@ -1,6 +1,5 @@
 # Runtime configurations
 ALPHATRION_LOG_LEVEL=INFO
-ALPHATRION_AUTO_CLEANUP=true
 
 # Database configurations
 ALPHATRION_METADATA_DB_URL=postgresql+psycopg2://alphatrion:alphatr1on@localhost:5432/alphatrion

.env.integration-test

@@ -2,7 +2,6 @@ ALPHATRION_METADATA_DB_URL=postgresql+psycopg2://alphatrion:alphatr1on@localhost
 ALPHATRION_ARTIFACT_REGISTRY_URL=localhost:5001
 ALPHATRION_ARTIFACT_INSECURE=true
 ALPHATRION_LOG_LEVEL=INFO
-ALPHATRION_AUTO_CLEANUP=true
 ALPHATRION_ENABLE_TRACING=true
 
 ALPHATRION_CLICKHOUSE_ENABLE_BATCH=true

.env.test

@@ -2,5 +2,4 @@ ALPHATRION_METADATA_INIT_TABLES=true
 ALPHATRION_METADATA_DB_URL=sqlite:///:memory:
 ALPHATRION_ARTIFACT_REGISTRY_URL=localhost:5001
 ALPHATRION_ARTIFACT_INSECURE=true
-ALPHATRION_LOG_LEVEL=INFO
-ALPHATRION_AUTO_CLEANUP=true
+ALPHATRION_LOG_LEVEL=INFO

alphatrion/agents/claude.py

@@ -245,6 +245,8 @@ def handle_stop():
         if not user_id and session:
             user_id = str(session.user_id)
 
+        user = metadb.get_user(user_id=uuid.UUID(user_id))
+
         if not team_id and session:
             team_id = str(session.team_id)
 
@@ -286,6 +288,7 @@ def handle_stop():
                 agent_id = metadb.create_agent(
                     name="claude",
                     type=AgentType.CLAUDE,
+                    org_id=user.org_id,
                     team_id=uuid.UUID(team_id),
                     user_id=uuid.UUID(user_id),
                 )
@@ -299,6 +302,7 @@ def handle_stop():
             db_session = metadb._session()
             session = AgentSession(
                 uuid=session_uuid,
+                org_id=user.org_id,
                 agent_id=agent_id,
                 team_id=uuid.UUID(team_id),
                 user_id=uuid.UUID(user_id),
@@ -589,6 +593,7 @@ def process_transcript_incremental(
             # Create run with aggregated tokens
             run_id = metadb.create_run(
                 session_id=uuid.UUID(session_id),
+                org_id=session.org_id,
                 team_id=session.team_id,
                 user_id=session.user_id,
                 status=run_status,

alphatrion/experiment/base.py

@@ -224,10 +224,8 @@ def _start(
 
             # reset to running status, also need to reset the tokens.
             if usage and "total_tokens" in usage:
-                # delete the tokens in the usage
-                usage.delete("total_tokens")
-                usage.delete("input_tokens")
-                usage.delete("output_tokens")
+                # delete the tokens in the usage - set to None instead of empty dict
+                usage = None
             self._runtime._metadb.update_experiment(
                 experiment_id=self._id,
                 status=Status.RUNNING,
@@ -241,6 +239,7 @@ def _start(
         else:
             self._id = self._runtime._metadb.create_experiment(
                 name=name,
+                org_id=self._runtime._org_id,
                 team_id=self._runtime._team_id,
                 user_id=self._runtime._user_id,
                 description=description,
@@ -262,6 +261,10 @@ def _start(
     def id(self) -> uuid.UUID:
         return self._id
 
+    @property
+    def team_id(self) -> uuid.UUID:
+        return self._runtime.team_id
+
     @property
     def config(self) -> ExperimentConfig:
         return self._config

alphatrion/log/log.py

@@ -114,7 +114,8 @@ async def log_metrics(metrics: dict[str, float]) -> bool:
         runtime._metadb.create_metric(
             key=key,
             value=value,
-            team_id=runtime._team_id,
+            org_id=runtime.org_id,
+            team_id=runtime.team_id,
             experiment_id=exp_id,
             run_id=run_id,
         )
@@ -183,6 +184,7 @@ async def log_dataset(
 
             id = runtime.metadb.create_dataset(
                 name=name,
+                org_id=runtime.org_id,
                 team_id=runtime.team_id,
                 user_id=runtime.user_id,
                 path=path,
@@ -198,6 +200,7 @@ async def log_dataset(
         )
         id = runtime.metadb.create_dataset(
             name=name,
+            org_id=runtime.org_id,
             team_id=runtime.team_id,
             user_id=runtime.user_id,
             path=path,

alphatrion/run/run.py

@@ -29,6 +29,7 @@ def _get_obj(self):
 
     def start(self, call_func: CallableEntry) -> None:
         self._id = self._runtime.metadb.create_run(
+            org_id=self._runtime.org_id,
             team_id=self._runtime.team_id,
             user_id=self._runtime.user_id,
             experiment_id=self._exp_id,

alphatrion/runtime/runtime.py

@@ -12,6 +12,7 @@
 def init(
     user_id: uuid.UUID,
     team_id: uuid.UUID | None = None,
+    org_id: uuid.UUID | None = None,
 ):
     """
     Initialize the AlphaTrion runtime environment.
@@ -20,13 +21,13 @@ def init(
         user_id: The user ID for the current user. You can generate a UUID
                  using `uuid.uuid4()`.
         team_id: The team ID for the current user. If not provided, will look
-                 for the first team
-                 associated with the user in the database.
+                 for the first team associated with the user in the database.
     """
     global __RUNTIME__
     __RUNTIME__ = Runtime(
         user_id=user_id,
         team_id=team_id,
+        org_id=org_id,
     )
 
 
@@ -42,6 +43,7 @@ class Runtime:
     __slots__ = (
         "_user_id",
         "_team_id",
+        "_org_id",
         "_metadb",
         "_tracestore",
         "_artifact",
@@ -53,6 +55,7 @@ def __init__(
         self,
         user_id: uuid.UUID,
         team_id: uuid.UUID | None = None,
+        org_id: uuid.UUID | None = None,
     ):
         storage_runtime.init()
         self._metadb = storage_runtime.storage_runtime().metadb
@@ -61,6 +64,7 @@ def __init__(
 
         self._user_id = user_id
         self._team_id = team_id
+        self._org_id = org_id
 
         if team_id is None:
             # If team_id is not provided, look for the first team associated with
@@ -73,6 +77,13 @@ def __init__(
                 )
             self._team_id = teams[0].uuid
 
+        if org_id is None:
+            # Look up org_id from user or team
+            user = self._metadb.get_user(user_id)
+            if not user:
+                raise ValueError(f"User {user_id} not found in the database.")
+            self._org_id = user.org_id
+
         self._root_path = os.getenv(envs.ROOT_PATH, os.path.expanduser("~/.alphatrion"))
         if not os.path.exists(self._root_path):
             os.makedirs(self._root_path, exist_ok=True)
@@ -97,6 +108,10 @@ def user_id(self) -> uuid.UUID:
     def team_id(self) -> uuid.UUID:
         return self._team_id
 
+    @property
+    def org_id(self) -> uuid.UUID:
+        return self._org_id
+
     @property
     def root_path(self) -> str:
         return self._root_path

alphatrion/server/cmd/app.py

@@ -8,6 +8,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from strawberry.fastapi import GraphQLRouter
 
+from alphatrion.server.graphql.context import get_context
 from alphatrion.server.graphql.schema import schema
 
 # Configure logging
@@ -99,8 +100,8 @@ async def receive():
     return response
 
 
-# Create GraphQL router
-graphql_app = GraphQLRouter(schema)
+# Create GraphQL router with context
+graphql_app = GraphQLRouter(schema, context_getter=get_context)
 
 # Mount /graphql endpoint
 app.include_router(graphql_app, prefix="/graphql")

alphatrion/server/cmd/main.py

@@ -98,6 +98,12 @@ def main():
         default="Default Team",
         help="Team name (default: Default Team)",
     )
+    init.add_argument(
+        "--org-name",
+        type=str,
+        default=None,
+        help="Organization name (auto-generated if not provided)",
+    )
     init.set_defaults(func=init_command)
 
     # run command (with subcommands)
@@ -173,32 +179,38 @@ def init_command(args):
         else f"{user_name.lower().replace(' ', '.')}@inftyai.com"
     )
     team_name = args.team_name
+    org_name = args.org_name if args.org_name else fake.company()
 
     try:
         metadb = runtime.storage_runtime().metadb
 
         console.print()
+        # Create organization
+        console.print(Text(f"🏢 Creating organization: {org_name}", style="bold cyan"))
+        org_id = metadb.create_organization(name=org_name)
         # Create user
         console.print(
             Text(f"👤 Creating user: {user_name} ({email})", style="bold cyan")
         )
-        user_id = metadb.create_user(username=user_name, email=email)
+        user_id = metadb.create_user(name=user_name, email=email, org_id=org_id)
 
         # Create team
         console.print(Text(f"🏢 Creating team: {team_name}", style="bold cyan"))
         team_id = metadb.create_team(
-            name=team_name, description=f"Team for {user_name}"
+            name=team_name, description=f"Team for {user_name}", org_id=org_id
         )
         # Add user to team
         metadb.add_user_to_team(user_id=user_id, team_id=team_id)
 
         console.print()
         console.print(Text("✅ Initialization successful!", style="bold green"))
         console.print()
-        console.print(Text("📋 Your user ID:", style="bold yellow"))
-        console.print(Text(f"   {user_id}", style="bold cyan"))
+        console.print(Text("📋 Your organization ID:", style="bold yellow"))
+        console.print(Text(f"   {org_id}", style="bold cyan"))
         console.print(Text("   Your team ID:", style="bold yellow"))
         console.print(Text(f"   {team_id}", style="bold cyan"))
+        console.print(Text("   Your user ID:", style="bold yellow"))
+        console.print(Text(f"   {user_id}", style="bold cyan"))
         console.print()
         console.print(
             Text(
@@ -296,6 +308,7 @@ def run_agent_command(args):
             agent_id = metadb.create_agent(
                 name=agent_name,
                 type=agent_type,
+                org_id=metadb.get_user(uuid.UUID(user_id)).org_id,
                 team_id=uuid.UUID(team_id),
                 user_id=uuid.UUID(user_id),
             )
@@ -529,12 +542,34 @@ def start_dashboard(args):
     # Create HTTP client for proxying requests to backend
     http_client = httpx.AsyncClient(base_url=args.backend_url, timeout=30.0)
 
-    # Endpoint to get current user ID (for frontend)
+    # Endpoint to get current user ID and org ID (for frontend)
     @app.get("/api/config")
     async def get_config():
+        import contextlib
+        import uuid
+
+        from alphatrion.storage import runtime as storage_runtime
+
+        # Initialize storage if not already done
+        with contextlib.suppress(Exception):
+            storage_runtime.init()
+
         config = {"userId": app.state.user_id}
+
+        # Look up user's org_id
+        try:
+            metadb = storage_runtime.storage_runtime().metadb
+            user = metadb.get_user(user_id=uuid.UUID(app.state.user_id))
+            if user:
+                config["orgId"] = str(user.org_id)
+        except Exception as e:
+            console.print(
+                Text(f"Warning: Could not fetch user org_id: {e}", style="yellow")
+            )
+
         if hasattr(app.state, "team_id"):
             config["teamId"] = app.state.team_id
+
         return config
 
     # Proxy /graphql requests to backend (MUST be before catch-all route)