From dd7e33faa1c0030abc73e147a5cd2bfdf57a697f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 2 Jan 2026 06:14:09 +0000 Subject: [PATCH 1/4] Initial plan From 840874c8ad217ca480740a8650211a9cba2bd1c6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 2 Jan 2026 06:24:30 +0000 Subject: [PATCH 2/4] Add customer-hosted agent with secure connection infrastructure Co-authored-by: ancient-kid <183126081+ancient-kid@users.noreply.github.com> --- backend/agent/README.md | 346 +++++++++++++++++ backend/agent/__init__.py | 26 ++ backend/agent/agent_config.example.json | 14 + backend/agent/auth.py | 267 ++++++++++++++ backend/agent/manager.py | 470 ++++++++++++++++++++++++ backend/agent/models.py | 175 +++++++++ backend/agent/run_agent.py | 364 ++++++++++++++++++ backend/agent/service.py | 407 ++++++++++++++++++++ backend/main.py | 323 ++++++++++++++++ my-app/prisma/schema.prisma | 29 ++ 10 files changed, 2421 insertions(+) create mode 100644 backend/agent/README.md create mode 100644 backend/agent/__init__.py create mode 100644 backend/agent/agent_config.example.json create mode 100644 backend/agent/auth.py create mode 100644 backend/agent/manager.py create mode 100644 backend/agent/models.py create mode 100644 backend/agent/run_agent.py create mode 100644 backend/agent/service.py diff --git a/backend/agent/README.md b/backend/agent/README.md new file mode 100644 index 0000000..7427d48 --- /dev/null +++ b/backend/agent/README.md @@ -0,0 +1,346 @@ +# Customer-Hosted Agent Setup Guide + +This guide explains how to deploy and configure the RELIX customer-hosted agent on your infrastructure. + +## Overview + +The customer-hosted agent allows you to securely connect your private databases to RELIX without exposing your database credentials or data to the cloud. The agent runs on your infrastructure and: + +- Receives authenticated query requests from the RELIX service +- Executes SQL queries locally against your database +- Returns only the query results (not raw credentials or full database access) +- Operates in read-only mode by default for safety + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ Your Infrastructure │ +│ │ +│ ┌──────────────────┐ ┌──────────────────┐ │ +│ │ Customer │ │ Your │ │ +│ │ Agent │─────▶│ Database │ │ +│ │ (Port 8443) │ │ (PostgreSQL) │ │ +│ └──────────────────┘ └──────────────────┘ │ +│ │ │ +│ │ HTTPS (TLS) │ +└───────────┼─────────────────────────────────────────────────────────┘ + │ + ▼ +┌───────────────────────────────────────────────────────────────────┐ +│ RELIX Cloud │ +│ ┌──────────────────┐ │ +│ │ RELIX API │◀──── Secure Token Authentication │ +│ │ Server │ │ +│ └──────────────────┘ │ +└───────────────────────────────────────────────────────────────────┘ +``` + +## Prerequisites + +- Python 3.9+ +- Access to your PostgreSQL database +- Network access from the agent to your database +- (Optional) SSL certificate for HTTPS + +## Quick Start + +### 1. Register Your Agent + +First, register a new agent through the RELIX API or dashboard: + +```bash +curl -X POST https://api.relix.com/agent/register \ + -H "Content-Type: application/json" \ + -d '{ + "agent_name": "Production DB Agent", + "host_url": "https://your-agent-host:8443", + "database_type": "postgres", + "allowed_schemas": ["public"], + "user_id": "your_user_id" + }' +``` + +Response: +```json +{ + "agent_id": "agent_abc123xyz", + "agent_token": "tok_xxxxxxxxxxxx", + "agent_secret": "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx", + "message": "Agent registered successfully", + "status": "pending" +} +``` + +**Important:** Save the `agent_id` and `agent_secret` securely. The secret is only shown once. + +### 2. Configure the Agent + +Create a configuration file `agent_config.json`: + +```json +{ + "agent_id": "agent_abc123xyz", + "agent_secret": "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx", + "database_url": "postgresql://user:password@localhost:5432/mydb", + "relix_server_url": "https://api.relix.com", + "allowed_schemas": ["public", "analytics"], + "read_only": true, + "max_rows": 10000, + "query_timeout": 30, + "heartbeat_interval": 30, + "port": 8443 +} +``` + +Or use environment variables: + +```bash +export AGENT_ID="agent_abc123xyz" +export AGENT_SECRET="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" +export DATABASE_URL="postgresql://user:password@localhost:5432/mydb" +export RELIX_SERVER_URL="https://api.relix.com" +export AGENT_PORT=8443 +export AGENT_ALLOWED_SCHEMAS="public,analytics" +export AGENT_READ_ONLY=true +``` + +### 3. Install Dependencies + +```bash +pip install fastapi uvicorn psycopg2-binary httpx pydantic +``` + +### 4. Run the Agent + +```bash +python -m agent.run_agent --config agent_config.json +``` + +Or with environment variables: + +```bash +python -m agent.run_agent +``` + +### 5. Verify Connection + +Check agent health: + +```bash +curl http://localhost:8443/health +``` + +Expected response: +```json +{ + "status": "healthy", + "agent_id": "agent_abc123xyz", + "version": "1.0.0", + "database_connected": true, + "timestamp": "2026-01-02T12:00:00Z" +} +``` + +## Configuration Options + +| Option | Description | Default | +|--------|-------------|---------| +| `agent_id` | Unique agent ID from registration | Required | +| `agent_secret` | Secret key from registration | Required | +| `database_url` | PostgreSQL connection string | Required | +| `relix_server_url` | RELIX API server URL | Required | +| `allowed_schemas` | Schemas the agent can query | `["public"]` | +| `read_only` | Only allow SELECT queries | `true` | +| `max_rows` | Maximum rows per query | `10000` | +| `query_timeout` | Query timeout in seconds | `30` | +| `heartbeat_interval` | Heartbeat interval in seconds | `30` | +| `port` | Port to listen on | `8443` | +| `ssl_cert` | Path to SSL certificate | `null` | +| `ssl_key` | Path to SSL private key | `null` | + +## Security Best Practices + +### 1. Use Read-Only Database User + +Create a dedicated read-only database user for the agent: + +```sql +CREATE USER relix_agent WITH PASSWORD 'secure_password'; +GRANT CONNECT ON DATABASE mydb TO relix_agent; +GRANT USAGE ON SCHEMA public TO relix_agent; +GRANT SELECT ON ALL TABLES IN SCHEMA public TO relix_agent; +ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO relix_agent; +``` + +### 2. Enable SSL/TLS + +For production, always use SSL: + +```json +{ + "ssl_cert": "/path/to/fullchain.pem", + "ssl_key": "/path/to/privkey.pem" +} +``` + +### 3. Network Security + +- Run the agent on a private network when possible +- Use a firewall to restrict incoming connections +- Only allow connections from RELIX IP addresses + +### 4. Keep Secrets Secure + +- Never commit `agent_config.json` with real credentials +- Use environment variables or secret management systems +- Rotate the agent secret periodically + +## Deployment Options + +### Docker + +```dockerfile +FROM python:3.11-slim + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY agent/ ./agent/ + +EXPOSE 8443 + +CMD ["python", "-m", "agent.run_agent"] +``` + +```bash +docker build -t relix-agent . +docker run -d \ + -p 8443:8443 \ + -e AGENT_ID=your_agent_id \ + -e AGENT_SECRET=your_secret \ + -e DATABASE_URL=postgresql://... \ + -e RELIX_SERVER_URL=https://api.relix.com \ + relix-agent +``` + +### Kubernetes + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: relix-agent +spec: + replicas: 1 + selector: + matchLabels: + app: relix-agent + template: + metadata: + labels: + app: relix-agent + spec: + containers: + - name: agent + image: relix-agent:latest + ports: + - containerPort: 8443 + envFrom: + - secretRef: + name: relix-agent-secrets +--- +apiVersion: v1 +kind: Secret +metadata: + name: relix-agent-secrets +type: Opaque +stringData: + AGENT_ID: "agent_xxx" + AGENT_SECRET: "xxx" + DATABASE_URL: "postgresql://..." + RELIX_SERVER_URL: "https://api.relix.com" +``` + +### systemd Service + +```ini +[Unit] +Description=RELIX Customer Agent +After=network.target postgresql.service + +[Service] +Type=simple +User=relix +WorkingDirectory=/opt/relix-agent +EnvironmentFile=/etc/relix-agent/env +ExecStart=/opt/relix-agent/venv/bin/python -m agent.run_agent +Restart=always +RestartSec=5 + +[Install] +WantedBy=multi-user.target +``` + +## Troubleshooting + +### Agent not connecting + +1. Check network connectivity to RELIX server +2. Verify agent credentials are correct +3. Check firewall rules +4. Review agent logs for errors + +### Database connection failed + +1. Verify `database_url` is correct +2. Check database user permissions +3. Ensure database is accessible from agent host + +### Queries timing out + +1. Increase `query_timeout` setting +2. Optimize slow queries +3. Check database performance + +### Authentication errors + +1. Verify `agent_id` and `agent_secret` match registration +2. Check if agent token has expired +3. Re-register the agent if needed + +## API Endpoints + +### Health Check + +``` +GET /health +``` + +Returns agent health status and database connectivity. + +### Status + +``` +GET /status +``` + +Returns detailed agent status including query statistics. + +### Query Execution + +``` +POST /query +``` + +Executes a query (called by RELIX server, not directly). + +## Support + +For issues or questions: + +1. Check the troubleshooting guide above +2. Review agent logs +3. Contact RELIX support diff --git a/backend/agent/__init__.py b/backend/agent/__init__.py new file mode 100644 index 0000000..fb24b57 --- /dev/null +++ b/backend/agent/__init__.py @@ -0,0 +1,26 @@ +""" +Customer-Hosted Agent Module + +This module provides the infrastructure for customer-hosted agents that can +securely connect to the main RELIX service. Agents run on customer infrastructure +and execute database queries locally, sending only results back to the main service. + +Components: +- auth: Token-based authentication for agent-service communication +- models: Pydantic models for agent configuration and messages +- service: Agent service that runs on customer infrastructure +""" + +from .models import AgentConfig, AgentRegistration, AgentStatus, AgentQueryRequest, AgentQueryResponse +from .auth import AgentAuthenticator, generate_agent_token, validate_agent_token + +__all__ = [ + "AgentConfig", + "AgentRegistration", + "AgentStatus", + "AgentQueryRequest", + "AgentQueryResponse", + "AgentAuthenticator", + "generate_agent_token", + "validate_agent_token", +] diff --git a/backend/agent/agent_config.example.json b/backend/agent/agent_config.example.json new file mode 100644 index 0000000..7315234 --- /dev/null +++ b/backend/agent/agent_config.example.json @@ -0,0 +1,14 @@ +{ + "agent_id": "YOUR_AGENT_ID", + "agent_secret": "YOUR_AGENT_SECRET", + "database_url": "postgresql://user:password@localhost:5432/your_database", + "relix_server_url": "https://api.relix.com", + "allowed_schemas": ["public"], + "read_only": true, + "max_rows": 10000, + "query_timeout": 30, + "heartbeat_interval": 30, + "port": 8443, + "ssl_cert": null, + "ssl_key": null +} diff --git a/backend/agent/auth.py b/backend/agent/auth.py new file mode 100644 index 0000000..eb44d7e --- /dev/null +++ b/backend/agent/auth.py @@ -0,0 +1,267 @@ +""" +Agent Authentication + +Secure token-based authentication for agent-service communication. +Uses HMAC-based tokens with time-limited validity. +""" + +import os +import hmac +import hashlib +import secrets +import base64 +import json +from datetime import datetime, timedelta +from typing import Optional, Tuple, Dict, Any + + +# Token validity period +TOKEN_VALIDITY_HOURS = 24 +SIGNATURE_ALGORITHM = "sha256" + + +def generate_agent_credentials() -> Tuple[str, str, str]: + """ + Generate a new set of agent credentials. + + Returns: + Tuple of (agent_id, agent_token, agent_secret) + """ + agent_id = f"agent_{secrets.token_hex(12)}" + agent_token = f"tok_{secrets.token_hex(24)}" + agent_secret = secrets.token_hex(32) + + return agent_id, agent_token, agent_secret + + +def generate_agent_token(agent_id: str, agent_secret: str, validity_hours: int = TOKEN_VALIDITY_HOURS) -> str: + """ + Generate a time-limited access token for agent authentication. + + Args: + agent_id: The agent's unique identifier + agent_secret: The agent's secret key + validity_hours: How long the token is valid + + Returns: + Base64-encoded signed token + """ + expiry = datetime.utcnow() + timedelta(hours=validity_hours) + + payload = { + "agent_id": agent_id, + "exp": expiry.isoformat(), + "iat": datetime.utcnow().isoformat(), + "nonce": secrets.token_hex(8) + } + + payload_json = json.dumps(payload, sort_keys=True) + payload_bytes = payload_json.encode('utf-8') + + # Create HMAC signature + signature = hmac.new( + agent_secret.encode('utf-8'), + payload_bytes, + hashlib.sha256 + ).digest() + + # Combine payload and signature + token_data = payload_bytes + b'.' + signature + + return base64.urlsafe_b64encode(token_data).decode('utf-8') + + +def validate_agent_token(token: str, agent_id: str, agent_secret: str) -> Tuple[bool, Optional[str]]: + """ + Validate an agent access token. + + Args: + token: The token to validate + agent_id: Expected agent ID + agent_secret: The agent's secret key + + Returns: + Tuple of (is_valid, error_message) + """ + try: + # Decode token + token_data = base64.urlsafe_b64decode(token.encode('utf-8')) + + # Split payload and signature + parts = token_data.rsplit(b'.', 1) + if len(parts) != 2: + return False, "Invalid token format" + + payload_bytes, received_signature = parts + + # Verify signature + expected_signature = hmac.new( + agent_secret.encode('utf-8'), + payload_bytes, + hashlib.sha256 + ).digest() + + if not hmac.compare_digest(received_signature, expected_signature): + return False, "Invalid signature" + + # Parse and validate payload + payload = json.loads(payload_bytes.decode('utf-8')) + + # Check agent ID + if payload.get("agent_id") != agent_id: + return False, "Agent ID mismatch" + + # Check expiry + expiry = datetime.fromisoformat(payload["exp"]) + if datetime.utcnow() > expiry: + return False, "Token expired" + + return True, None + + except Exception as e: + return False, f"Token validation error: {str(e)}" + + +def sign_payload(payload: Dict[str, Any], agent_secret: str) -> str: + """ + Sign a payload with the agent secret. + + Args: + payload: Dictionary payload to sign + agent_secret: The agent's secret key + + Returns: + Base64-encoded signature + """ + payload_json = json.dumps(payload, sort_keys=True, default=str) + + signature = hmac.new( + agent_secret.encode('utf-8'), + payload_json.encode('utf-8'), + hashlib.sha256 + ).digest() + + return base64.urlsafe_b64encode(signature).decode('utf-8') + + +def verify_payload_signature(payload: Dict[str, Any], signature: str, agent_secret: str) -> bool: + """ + Verify a payload signature. + + Args: + payload: Dictionary payload that was signed + signature: Base64-encoded signature to verify + agent_secret: The agent's secret key + + Returns: + True if signature is valid + """ + try: + expected_signature = sign_payload(payload, agent_secret) + return hmac.compare_digest(signature, expected_signature) + except Exception: + return False + + +class AgentAuthenticator: + """ + Handles agent authentication and message signing. + """ + + def __init__(self, agent_id: str, agent_secret: str): + self.agent_id = agent_id + self.agent_secret = agent_secret + self._current_token: Optional[str] = None + self._token_expiry: Optional[datetime] = None + + def get_token(self, force_refresh: bool = False) -> str: + """ + Get a valid access token, refreshing if necessary. + + Args: + force_refresh: Force generation of a new token + + Returns: + Valid access token + """ + # Check if current token is still valid + if not force_refresh and self._current_token and self._token_expiry: + # Refresh if less than 1 hour remaining + if datetime.utcnow() < self._token_expiry - timedelta(hours=1): + return self._current_token + + # Generate new token + self._current_token = generate_agent_token(self.agent_id, self.agent_secret) + self._token_expiry = datetime.utcnow() + timedelta(hours=TOKEN_VALIDITY_HOURS) + + return self._current_token + + def sign_request(self, payload: Dict[str, Any]) -> Dict[str, Any]: + """ + Sign a request payload for secure transmission. + + Args: + payload: Request payload + + Returns: + Payload with added authentication headers + """ + timestamp = datetime.utcnow().isoformat() + + # Add metadata to payload + signed_payload = { + **payload, + "_agent_id": self.agent_id, + "_timestamp": timestamp, + "_nonce": secrets.token_hex(8) + } + + # Generate signature + signature = sign_payload(signed_payload, self.agent_secret) + + return { + "payload": signed_payload, + "signature": signature, + "token": self.get_token() + } + + def verify_request(self, request_data: Dict[str, Any]) -> Tuple[bool, Optional[str], Optional[Dict[str, Any]]]: + """ + Verify an incoming signed request. + + Args: + request_data: The complete request with payload, signature, and token + + Returns: + Tuple of (is_valid, error_message, payload) + """ + try: + payload = request_data.get("payload") + signature = request_data.get("signature") + token = request_data.get("token") + + if not all([payload, signature, token]): + return False, "Missing required fields", None + + # Verify token + token_valid, token_error = validate_agent_token(token, self.agent_id, self.agent_secret) + if not token_valid: + return False, token_error, None + + # Verify signature + if not verify_payload_signature(payload, signature, self.agent_secret): + return False, "Invalid signature", None + + # Verify timestamp (within 5 minutes) + timestamp = datetime.fromisoformat(payload.get("_timestamp", "1970-01-01")) + if abs((datetime.utcnow() - timestamp).total_seconds()) > 300: + return False, "Request timestamp too old", None + + # Verify agent ID + if payload.get("_agent_id") != self.agent_id: + return False, "Agent ID mismatch", None + + return True, None, payload + + except Exception as e: + return False, f"Verification error: {str(e)}", None diff --git a/backend/agent/manager.py b/backend/agent/manager.py new file mode 100644 index 0000000..9015d8e --- /dev/null +++ b/backend/agent/manager.py @@ -0,0 +1,470 @@ +""" +Agent Manager + +Manages customer-hosted agents from the main RELIX service. +Handles registration, heartbeats, connection status, and query routing. +""" + +import os +import json +import asyncio +import logging +import httpx +from typing import Optional, Dict, Any, List, Tuple +from datetime import datetime, timedelta + +from .models import ( + AgentConfig, + AgentRegistration, + AgentRegistrationResponse, + AgentStatus, + AgentQueryRequest, + AgentQueryResponse, + AgentHeartbeat, + AgentStatusResponse +) +from .auth import generate_agent_credentials, AgentAuthenticator, sign_payload + +logger = logging.getLogger("agent.manager") + + +class AgentManager: + """ + Manages customer-hosted agents from the server side. + + Responsibilities: + - Register new agents and generate credentials + - Track agent status via heartbeats + - Route queries to appropriate agents + - Handle agent disconnections and failovers + """ + + def __init__(self, db_cursor_factory): + """ + Initialize the agent manager. + + Args: + db_cursor_factory: Function that returns a database cursor context manager + """ + self.db_cursor = db_cursor_factory + self._agent_cache: Dict[str, Dict[str, Any]] = {} + self._heartbeat_timestamps: Dict[str, datetime] = {} + + # Configuration + self.heartbeat_timeout_seconds = 60 # Agent considered disconnected after this + self.query_timeout_seconds = 30 + self.max_retries = 2 + + async def register_agent(self, registration: AgentRegistration) -> AgentRegistrationResponse: + """ + Register a new customer-hosted agent. + + Args: + registration: Agent registration details + + Returns: + Registration response with credentials + """ + # Generate credentials + agent_id, agent_token, agent_secret = generate_agent_credentials() + + # Store in database + config = AgentConfig( + agent_id=agent_id, + agent_name=registration.agent_name, + host_url=registration.host_url, + database_type=registration.database_type, + allowed_schemas=registration.allowed_schemas + ) + + with self.db_cursor(commit=True) as cur: + cur.execute( + """ + INSERT INTO "Agent" ( + id, "userId", name, "hostUrl", "databaseType", + "allowedSchemas", "agentToken", "agentSecret", + status, "createdAt", "updatedAt" + ) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW()) + """, + ( + agent_id, + registration.user_id, + registration.agent_name, + registration.host_url, + registration.database_type, + json.dumps(registration.allowed_schemas), + agent_token, # Store hashed in production + agent_secret, # Store hashed in production + AgentStatus.PENDING.value + ) + ) + + logger.info(f"Registered new agent: {agent_id} for user: {registration.user_id}") + + return AgentRegistrationResponse( + agent_id=agent_id, + agent_token=agent_token, + agent_secret=agent_secret, + message="Agent registered successfully. Use these credentials in your agent configuration.", + status=AgentStatus.PENDING + ) + + async def process_heartbeat(self, heartbeat: AgentHeartbeat) -> Dict[str, Any]: + """ + Process a heartbeat from an agent. + + Args: + heartbeat: Heartbeat message from agent + + Returns: + Acknowledgment response + """ + agent_id = heartbeat.agent_id + + # Validate token + agent_data = await self._get_agent_data(agent_id) + if not agent_data: + return {"success": False, "error": "Unknown agent"} + + # Update status + with self.db_cursor(commit=True) as cur: + cur.execute( + """ + UPDATE "Agent" + SET status = %s, "lastHeartbeat" = NOW(), "updatedAt" = NOW(), + "connectedDatabases" = %s, "activeConnections" = %s + WHERE id = %s + """, + ( + AgentStatus.CONNECTED.value, + json.dumps(heartbeat.connected_databases), + heartbeat.active_connections, + agent_id + ) + ) + + self._heartbeat_timestamps[agent_id] = datetime.utcnow() + + # Invalidate cache + if agent_id in self._agent_cache: + del self._agent_cache[agent_id] + + logger.debug(f"Heartbeat received from agent: {agent_id}") + + return {"success": True, "message": "Heartbeat acknowledged"} + + async def execute_query( + self, + agent_id: str, + sql: str, + parameters: Optional[Dict[str, Any]] = None, + schema_name: str = "public", + timeout: int = 30 + ) -> AgentQueryResponse: + """ + Execute a query through a customer-hosted agent. + + Args: + agent_id: Target agent ID + sql: SQL query to execute + parameters: Query parameters + schema_name: Target schema + timeout: Query timeout in seconds + + Returns: + Query response from agent + """ + import uuid + + # Get agent data + agent_data = await self._get_agent_data(agent_id) + if not agent_data: + return AgentQueryResponse( + query_id="error", + success=False, + error="Agent not found" + ) + + # Check agent status + if agent_data.get("status") != AgentStatus.CONNECTED.value: + return AgentQueryResponse( + query_id="error", + success=False, + error=f"Agent not connected. Status: {agent_data.get('status')}" + ) + + # Build query request + query_id = f"query_{uuid.uuid4().hex[:12]}" + request = AgentQueryRequest( + query_id=query_id, + sql=sql, + parameters=parameters, + schema_name=schema_name, + timeout_seconds=min(timeout, self.query_timeout_seconds) + ) + + # Sign request + authenticator = AgentAuthenticator(agent_id, agent_data["agentSecret"]) + signed_request = authenticator.sign_request(request.model_dump()) + + # Send to agent + host_url = agent_data["hostUrl"].rstrip("/") + + try: + async with httpx.AsyncClient(timeout=timeout + 5) as client: + response = await client.post( + f"{host_url}/query", + json=signed_request, + headers={"Content-Type": "application/json"} + ) + + if response.status_code != 200: + return AgentQueryResponse( + query_id=query_id, + success=False, + error=f"Agent returned status {response.status_code}" + ) + + result = response.json() + return AgentQueryResponse(**result) + + except httpx.TimeoutException: + logger.error(f"Query timeout for agent {agent_id}") + return AgentQueryResponse( + query_id=query_id, + success=False, + error="Query timed out" + ) + except Exception as e: + logger.error(f"Query execution error: {e}") + return AgentQueryResponse( + query_id=query_id, + success=False, + error=f"Communication error: {str(e)}" + ) + + async def get_agent_status(self, agent_id: str, user_id: str) -> Optional[AgentStatusResponse]: + """ + Get status information for an agent. + + Args: + agent_id: Agent ID + user_id: User ID for authorization + + Returns: + Agent status or None if not found + """ + with self.db_cursor() as cur: + cur.execute( + """ + SELECT * FROM "Agent" WHERE id = %s AND "userId" = %s + """, + (agent_id, user_id) + ) + agent = cur.fetchone() + + if not agent: + return None + + # Calculate uptime + uptime = None + last_heartbeat = agent.get("lastHeartbeat") + if last_heartbeat: + uptime = int((datetime.utcnow() - last_heartbeat).total_seconds()) + + return AgentStatusResponse( + agent_id=agent["id"], + agent_name=agent["name"], + status=AgentStatus(agent["status"]), + host_url=agent["hostUrl"], + database_type=agent["databaseType"], + last_heartbeat=last_heartbeat, + uptime_seconds=uptime, + queries_executed=agent.get("queriesExecuted", 0), + avg_response_time_ms=agent.get("avgResponseTime", 0), + error_rate=agent.get("errorRate", 0) + ) + + async def list_agents(self, user_id: str) -> List[AgentStatusResponse]: + """ + List all agents for a user. + + Args: + user_id: User ID + + Returns: + List of agent status responses + """ + with self.db_cursor() as cur: + cur.execute( + """ + SELECT * FROM "Agent" WHERE "userId" = %s ORDER BY "createdAt" DESC + """, + (user_id,) + ) + agents = cur.fetchall() + + result = [] + for agent in agents: + uptime = None + last_heartbeat = agent.get("lastHeartbeat") + if last_heartbeat: + uptime = int((datetime.utcnow() - last_heartbeat).total_seconds()) + + result.append(AgentStatusResponse( + agent_id=agent["id"], + agent_name=agent["name"], + status=AgentStatus(agent["status"]), + host_url=agent["hostUrl"], + database_type=agent["databaseType"], + last_heartbeat=last_heartbeat, + uptime_seconds=uptime, + queries_executed=agent.get("queriesExecuted", 0), + avg_response_time_ms=agent.get("avgResponseTime", 0), + error_rate=agent.get("errorRate", 0) + )) + + return result + + async def delete_agent(self, agent_id: str, user_id: str) -> bool: + """ + Delete an agent. + + Args: + agent_id: Agent ID + user_id: User ID for authorization + + Returns: + True if deleted, False if not found + """ + with self.db_cursor(commit=True) as cur: + cur.execute( + """ + DELETE FROM "Agent" WHERE id = %s AND "userId" = %s + """, + (agent_id, user_id) + ) + deleted = cur.rowcount > 0 + + if deleted: + # Clean up cache + if agent_id in self._agent_cache: + del self._agent_cache[agent_id] + if agent_id in self._heartbeat_timestamps: + del self._heartbeat_timestamps[agent_id] + + logger.info(f"Deleted agent: {agent_id}") + + return deleted + + async def test_agent_connection(self, agent_id: str, user_id: str) -> Tuple[bool, str, float]: + """ + Test connection to an agent. + + Args: + agent_id: Agent ID + user_id: User ID for authorization + + Returns: + Tuple of (success, message, latency_ms) + """ + import time + + agent_data = await self._get_agent_data(agent_id) + if not agent_data: + return False, "Agent not found", 0 + + if agent_data.get("userId") != user_id: + return False, "Unauthorized", 0 + + host_url = agent_data["hostUrl"].rstrip("/") + + start_time = time.time() + try: + async with httpx.AsyncClient(timeout=10) as client: + response = await client.get(f"{host_url}/health") + latency = (time.time() - start_time) * 1000 + + if response.status_code == 200: + return True, "Connection successful", latency + else: + return False, f"Agent returned status {response.status_code}", latency + + except httpx.TimeoutException: + return False, "Connection timed out", (time.time() - start_time) * 1000 + except Exception as e: + return False, f"Connection error: {str(e)}", (time.time() - start_time) * 1000 + + async def _get_agent_data(self, agent_id: str) -> Optional[Dict[str, Any]]: + """Get agent data from cache or database.""" + # Check cache + if agent_id in self._agent_cache: + cached = self._agent_cache[agent_id] + if datetime.utcnow() - cached["_cached_at"] < timedelta(minutes=5): + return cached + + # Fetch from database + with self.db_cursor() as cur: + cur.execute( + """SELECT * FROM "Agent" WHERE id = %s""", + (agent_id,) + ) + agent = cur.fetchone() + + if agent: + agent_dict = dict(agent) + agent_dict["_cached_at"] = datetime.utcnow() + self._agent_cache[agent_id] = agent_dict + return agent_dict + + return None + + async def check_stale_agents(self) -> List[str]: + """ + Check for agents that haven't sent heartbeats recently. + Updates their status to disconnected. + + Returns: + List of agent IDs marked as disconnected + """ + threshold = datetime.utcnow() - timedelta(seconds=self.heartbeat_timeout_seconds) + + with self.db_cursor(commit=True) as cur: + cur.execute( + """ + UPDATE "Agent" + SET status = %s, "updatedAt" = NOW() + WHERE status = %s AND "lastHeartbeat" < %s + RETURNING id + """, + (AgentStatus.DISCONNECTED.value, AgentStatus.CONNECTED.value, threshold) + ) + stale_agents = [row["id"] for row in cur.fetchall()] + + # Clear cache for stale agents + for agent_id in stale_agents: + if agent_id in self._agent_cache: + del self._agent_cache[agent_id] + logger.warning(f"Agent marked as disconnected: {agent_id}") + + return stale_agents + + +# Global instance (will be initialized in main.py) +_agent_manager: Optional[AgentManager] = None + + +def get_agent_manager() -> AgentManager: + """Get the global agent manager instance.""" + global _agent_manager + if _agent_manager is None: + raise RuntimeError("Agent manager not initialized") + return _agent_manager + + +def init_agent_manager(db_cursor_factory) -> AgentManager: + """Initialize the global agent manager.""" + global _agent_manager + _agent_manager = AgentManager(db_cursor_factory) + return _agent_manager diff --git a/backend/agent/models.py b/backend/agent/models.py new file mode 100644 index 0000000..650c2ef --- /dev/null +++ b/backend/agent/models.py @@ -0,0 +1,175 @@ +""" +Agent Models + +Pydantic models for customer-hosted agent configuration, registration, +status tracking, and query message payloads. +""" + +from pydantic import BaseModel, Field +from typing import Optional, List, Dict, Any +from datetime import datetime +from enum import Enum + + +class AgentStatus(str, Enum): + """Agent connection status""" + PENDING = "pending" + CONNECTED = "connected" + DISCONNECTED = "disconnected" + ERROR = "error" + + +class AgentConfig(BaseModel): + """Configuration for a customer-hosted agent""" + agent_id: str = Field(..., description="Unique agent identifier") + agent_name: str = Field(..., description="Human-readable agent name") + host_url: str = Field(..., description="URL where agent is hosted") + database_type: str = Field(default="postgres", description="Database type (postgres, mysql, etc.)") + allowed_schemas: List[str] = Field(default=["public"], description="Schemas agent can access") + max_rows_per_query: int = Field(default=10000, description="Maximum rows returned per query") + timeout_seconds: int = Field(default=30, description="Query timeout in seconds") + read_only: bool = Field(default=True, description="Only allow SELECT queries") + + class Config: + json_schema_extra = { + "example": { + "agent_id": "agent_abc123", + "agent_name": "Production DB Agent", + "host_url": "https://agent.customer.com:8443", + "database_type": "postgres", + "allowed_schemas": ["public", "analytics"], + "max_rows_per_query": 10000, + "timeout_seconds": 30, + "read_only": True + } + } + + +class AgentRegistration(BaseModel): + """Request payload for registering a new agent""" + agent_name: str = Field(..., description="Human-readable agent name") + host_url: str = Field(..., description="URL where agent will be hosted") + database_type: str = Field(default="postgres", description="Database type") + allowed_schemas: List[str] = Field(default=["public"], description="Schemas agent can access") + user_id: str = Field(..., description="Owner user ID") + + class Config: + json_schema_extra = { + "example": { + "agent_name": "My Production Agent", + "host_url": "https://agent.mycompany.com:8443", + "database_type": "postgres", + "allowed_schemas": ["public"], + "user_id": "user_clerk_xxx" + } + } + + +class AgentRegistrationResponse(BaseModel): + """Response after successful agent registration""" + agent_id: str + agent_token: str + agent_secret: str + message: str + status: AgentStatus + + +class AgentHeartbeat(BaseModel): + """Heartbeat message from agent to server""" + agent_id: str + agent_token: str + timestamp: datetime = Field(default_factory=datetime.utcnow) + status: str = Field(default="healthy") + connected_databases: List[str] = Field(default=[]) + active_connections: int = Field(default=0) + + class Config: + json_schema_extra = { + "example": { + "agent_id": "agent_abc123", + "agent_token": "tok_xxx", + "timestamp": "2026-01-02T12:00:00Z", + "status": "healthy", + "connected_databases": ["production"], + "active_connections": 5 + } + } + + +class AgentQueryRequest(BaseModel): + """Query request sent to agent for execution""" + query_id: str = Field(..., description="Unique query identifier") + sql: str = Field(..., description="SQL query to execute") + parameters: Optional[Dict[str, Any]] = Field(default=None, description="Query parameters") + schema_name: Optional[str] = Field(default="public", description="Target schema") + timeout_seconds: int = Field(default=30, description="Query timeout") + max_rows: int = Field(default=10000, description="Maximum rows to return") + + class Config: + json_schema_extra = { + "example": { + "query_id": "query_xyz789", + "sql": "SELECT * FROM users WHERE created_at > $1 LIMIT 100", + "parameters": {"$1": "2025-01-01"}, + "schema_name": "public", + "timeout_seconds": 30, + "max_rows": 100 + } + } + + +class AgentQueryResponse(BaseModel): + """Query response from agent""" + query_id: str + success: bool + rows: Optional[List[Dict[str, Any]]] = None + row_count: int = 0 + columns: Optional[List[str]] = None + column_types: Optional[Dict[str, str]] = None + execution_time_ms: float = 0 + error: Optional[str] = None + truncated: bool = False + + class Config: + json_schema_extra = { + "example": { + "query_id": "query_xyz789", + "success": True, + "rows": [{"id": 1, "name": "John"}], + "row_count": 1, + "columns": ["id", "name"], + "column_types": {"id": "integer", "name": "varchar"}, + "execution_time_ms": 45.2, + "error": None, + "truncated": False + } + } + + +class AgentStatusResponse(BaseModel): + """Agent status information""" + agent_id: str + agent_name: str + status: AgentStatus + host_url: str + database_type: str + last_heartbeat: Optional[datetime] = None + uptime_seconds: Optional[int] = None + queries_executed: int = 0 + avg_response_time_ms: float = 0 + error_rate: float = 0 + + +class AgentConnectionTest(BaseModel): + """Request to test agent connection""" + agent_id: str + agent_token: str + + +class AgentConnectionTestResponse(BaseModel): + """Response from agent connection test""" + success: bool + latency_ms: float + message: str + agent_version: Optional[str] = None + database_connected: bool = False diff --git a/backend/agent/run_agent.py b/backend/agent/run_agent.py new file mode 100644 index 0000000..32867b0 --- /dev/null +++ b/backend/agent/run_agent.py @@ -0,0 +1,364 @@ +#!/usr/bin/env python3 +""" +Customer-Hosted Agent Runner + +This script runs the customer-hosted agent as a standalone service. +Deploy this on your infrastructure to securely connect your databases +to the RELIX NL-to-SQL service. + +Usage: + python run_agent.py --config agent_config.json + + Or set environment variables: + AGENT_ID=agent_xxx + AGENT_SECRET=xxx + DATABASE_URL=postgresql://... + RELIX_SERVER_URL=https://api.relix.com + + python run_agent.py + +Configuration file format (agent_config.json): + { + "agent_id": "agent_xxx", + "agent_secret": "your_agent_secret", + "database_url": "postgresql://user:pass@localhost:5432/mydb", + "relix_server_url": "https://api.relix.com", + "allowed_schemas": ["public", "analytics"], + "read_only": true, + "max_rows": 10000, + "query_timeout": 30, + "heartbeat_interval": 30, + "port": 8443, + "ssl_cert": "/path/to/cert.pem", + "ssl_key": "/path/to/key.pem" + } +""" + +import os +import sys +import json +import argparse +import asyncio +import logging +import signal +from datetime import datetime +from typing import Optional + +# Add parent directory to path for imports +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from fastapi import FastAPI, HTTPException, Request +from fastapi.middleware.cors import CORSMiddleware +import uvicorn +import httpx + +from agent.service import AgentService +from agent.models import AgentQueryRequest, AgentQueryResponse +from agent.auth import validate_agent_token + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger("agent.runner") + +# Global agent service instance +agent_service: Optional[AgentService] = None +config: dict = {} + + +def create_app() -> FastAPI: + """Create the FastAPI application for the agent.""" + app = FastAPI( + title="RELIX Customer Agent", + description="Customer-hosted agent for secure database access", + version="1.0.0" + ) + + # Configure CORS - only allow RELIX server + relix_url = config.get("relix_server_url", "") + app.add_middleware( + CORSMiddleware, + allow_origins=[relix_url] if relix_url else ["*"], + allow_credentials=True, + allow_methods=["GET", "POST"], + allow_headers=["*"], + ) + + @app.get("/health") + async def health(): + """Health check endpoint.""" + if agent_service is None: + return {"status": "not_initialized"} + + db_ok, db_msg = agent_service.test_database_connection() + + return { + "status": "healthy" if db_ok else "degraded", + "agent_id": agent_service.agent_id, + "version": AgentService.VERSION, + "database_connected": db_ok, + "database_message": db_msg, + "timestamp": datetime.utcnow().isoformat() + } + + @app.get("/status") + async def status(): + """Get detailed agent status.""" + if agent_service is None: + raise HTTPException(status_code=503, detail="Agent not initialized") + + return agent_service.get_status() + + @app.post("/query") + async def execute_query(request: Request): + """ + Execute a query from the RELIX server. + + Request body is a signed payload from the server. + """ + if agent_service is None: + raise HTTPException(status_code=503, detail="Agent not initialized") + + try: + body = await request.json() + + # Extract and verify token + token = body.get("token") + if not token: + raise HTTPException(status_code=401, detail="Missing authentication token") + + is_valid, error = agent_service.verify_request(token) + if not is_valid: + logger.warning(f"Invalid request token: {error}") + raise HTTPException(status_code=401, detail=f"Authentication failed: {error}") + + # Extract payload + payload = body.get("payload", {}) + + # Build query request + query_request = AgentQueryRequest( + query_id=payload.get("query_id", "unknown"), + sql=payload.get("sql", ""), + parameters=payload.get("parameters"), + schema_name=payload.get("schema_name", "public"), + timeout_seconds=payload.get("timeout_seconds", 30), + max_rows=payload.get("max_rows", 10000) + ) + + # Execute query + result = agent_service.execute_query(query_request) + + return result.model_dump() + + except HTTPException: + raise + except Exception as e: + logger.error(f"Query execution error: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + return app + + +async def send_heartbeat(): + """Send periodic heartbeats to the RELIX server.""" + global agent_service, config + + relix_url = config.get("relix_server_url", "").rstrip("/") + interval = config.get("heartbeat_interval", 30) + + if not relix_url: + logger.warning("No RELIX server URL configured, heartbeats disabled") + return + + logger.info(f"Starting heartbeat sender (interval: {interval}s)") + + while True: + try: + if agent_service: + heartbeat = agent_service.get_heartbeat() + + async with httpx.AsyncClient(timeout=10) as client: + response = await client.post( + f"{relix_url}/agent/heartbeat", + json=heartbeat.model_dump(), + headers={"Content-Type": "application/json"} + ) + + if response.status_code == 200: + logger.debug("Heartbeat sent successfully") + else: + logger.warning(f"Heartbeat failed: {response.status_code}") + + await asyncio.sleep(interval) + + except asyncio.CancelledError: + logger.info("Heartbeat sender stopped") + break + except Exception as e: + logger.error(f"Heartbeat error: {e}") + await asyncio.sleep(interval) + + +def load_config(config_path: Optional[str] = None) -> dict: + """Load configuration from file or environment variables.""" + config = {} + + # Try loading from file + if config_path and os.path.exists(config_path): + with open(config_path, 'r') as f: + config = json.load(f) + logger.info(f"Loaded configuration from {config_path}") + + # Override with environment variables + env_mapping = { + "AGENT_ID": "agent_id", + "AGENT_SECRET": "agent_secret", + "DATABASE_URL": "database_url", + "RELIX_SERVER_URL": "relix_server_url", + "AGENT_PORT": "port", + "AGENT_ALLOWED_SCHEMAS": "allowed_schemas", + "AGENT_READ_ONLY": "read_only", + "AGENT_MAX_ROWS": "max_rows", + "AGENT_QUERY_TIMEOUT": "query_timeout", + "AGENT_HEARTBEAT_INTERVAL": "heartbeat_interval", + "AGENT_SSL_CERT": "ssl_cert", + "AGENT_SSL_KEY": "ssl_key", + } + + for env_var, config_key in env_mapping.items(): + value = os.getenv(env_var) + if value is not None: + # Type conversion + if config_key in ["port", "max_rows", "query_timeout", "heartbeat_interval"]: + value = int(value) + elif config_key == "read_only": + value = value.lower() in ("true", "1", "yes") + elif config_key == "allowed_schemas": + value = [s.strip() for s in value.split(",")] + + config[config_key] = value + + return config + + +def validate_config(config: dict) -> bool: + """Validate required configuration.""" + required = ["agent_id", "agent_secret", "database_url"] + missing = [key for key in required if not config.get(key)] + + if missing: + logger.error(f"Missing required configuration: {', '.join(missing)}") + return False + + return True + + +def main(): + """Main entry point.""" + global agent_service, config + + parser = argparse.ArgumentParser(description="Run RELIX Customer Agent") + parser.add_argument( + "--config", "-c", + help="Path to configuration file", + default="agent_config.json" + ) + parser.add_argument( + "--port", "-p", + type=int, + help="Port to listen on (default: 8443)", + default=None + ) + parser.add_argument( + "--host", + help="Host to bind to (default: 0.0.0.0)", + default="0.0.0.0" + ) + + args = parser.parse_args() + + # Load configuration + config = load_config(args.config) + + # Override port from command line + if args.port: + config["port"] = args.port + + # Validate configuration + if not validate_config(config): + sys.exit(1) + + # Initialize agent service + agent_service = AgentService( + agent_id=config["agent_id"], + agent_secret=config["agent_secret"], + database_url=config["database_url"], + allowed_schemas=config.get("allowed_schemas", ["public"]), + read_only=config.get("read_only", True), + max_rows=config.get("max_rows", 10000), + query_timeout=config.get("query_timeout", 30) + ) + + # Test database connection + db_ok, db_msg = agent_service.test_database_connection() + if not db_ok: + logger.error(f"Database connection failed: {db_msg}") + sys.exit(1) + + logger.info(f"Database connection successful") + + # Start agent + agent_service.start() + + # Create FastAPI app + app = create_app() + + # Get SSL configuration + ssl_cert = config.get("ssl_cert") + ssl_key = config.get("ssl_key") + + ssl_kwargs = {} + if ssl_cert and ssl_key: + if os.path.exists(ssl_cert) and os.path.exists(ssl_key): + ssl_kwargs = { + "ssl_certfile": ssl_cert, + "ssl_keyfile": ssl_key + } + logger.info("SSL enabled") + else: + logger.warning("SSL certificate/key files not found, running without SSL") + + port = config.get("port", 8443) + + logger.info(f"Starting agent server on {args.host}:{port}") + logger.info(f"Agent ID: {config['agent_id']}") + + # Start heartbeat task + async def startup(): + asyncio.create_task(send_heartbeat()) + + app.add_event_handler("startup", startup) + + # Handle shutdown + def shutdown_handler(signum, frame): + logger.info("Shutting down...") + if agent_service: + agent_service.stop() + sys.exit(0) + + signal.signal(signal.SIGINT, shutdown_handler) + signal.signal(signal.SIGTERM, shutdown_handler) + + # Run server + uvicorn.run( + app, + host=args.host, + port=port, + **ssl_kwargs + ) + + +if __name__ == "__main__": + main() diff --git a/backend/agent/service.py b/backend/agent/service.py new file mode 100644 index 0000000..81e7574 --- /dev/null +++ b/backend/agent/service.py @@ -0,0 +1,407 @@ +""" +Customer-Hosted Agent Service + +This is the lightweight agent service that runs on customer infrastructure. +It receives queries from the main RELIX service and executes them against +the local database, returning only the results. + +Usage: + python -m agent.service --config agent_config.json + +Or programmatically: + from agent.service import AgentService + + service = AgentService(config) + service.start() +""" + +import os +import json +import time +import asyncio +import logging +from typing import Optional, Dict, Any, List +from datetime import datetime + +import psycopg2 +from psycopg2.extras import RealDictCursor + +from .models import ( + AgentConfig, + AgentQueryRequest, + AgentQueryResponse, + AgentHeartbeat, + AgentStatus +) +from .auth import AgentAuthenticator, validate_agent_token + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("agent.service") + + +class AgentService: + """ + Customer-hosted agent service for secure database query execution. + + This service: + 1. Connects to the customer's local database + 2. Receives authenticated query requests from RELIX service + 3. Executes queries locally (read-only by default) + 4. Returns results securely to the main service + 5. Sends periodic heartbeats to maintain connection status + """ + + VERSION = "1.0.0" + + def __init__( + self, + agent_id: str, + agent_secret: str, + database_url: str, + allowed_schemas: List[str] = None, + read_only: bool = True, + max_rows: int = 10000, + query_timeout: int = 30 + ): + """ + Initialize the agent service. + + Args: + agent_id: Unique agent identifier from registration + agent_secret: Secret key from registration + database_url: PostgreSQL connection string for local database + allowed_schemas: List of schemas the agent can query + read_only: Only allow SELECT queries (default: True) + max_rows: Maximum rows to return per query (default: 10000) + query_timeout: Query timeout in seconds (default: 30) + """ + self.agent_id = agent_id + self.agent_secret = agent_secret + self.database_url = database_url + self.allowed_schemas = allowed_schemas or ["public"] + self.read_only = read_only + self.max_rows = max_rows + self.query_timeout = query_timeout + + self.authenticator = AgentAuthenticator(agent_id, agent_secret) + self._db_connection: Optional[psycopg2.extensions.connection] = None + self._is_running = False + self._queries_executed = 0 + self._total_execution_time = 0.0 + self._errors = 0 + self._start_time: Optional[datetime] = None + + logger.info(f"Agent service initialized: {agent_id}") + + def _get_db_connection(self) -> psycopg2.extensions.connection: + """Get or create database connection.""" + if self._db_connection is None or self._db_connection.closed: + self._db_connection = psycopg2.connect( + self.database_url, + cursor_factory=RealDictCursor + ) + if self.read_only: + self._db_connection.set_session(readonly=True) + logger.info("Database connection established") + return self._db_connection + + def _validate_query(self, sql: str) -> tuple[bool, Optional[str]]: + """ + Validate SQL query for safety. + + Args: + sql: SQL query string + + Returns: + Tuple of (is_valid, error_message) + """ + sql_upper = sql.strip().upper() + + # Check for read-only mode + if self.read_only: + allowed_prefixes = ("SELECT", "WITH", "EXPLAIN") + if not any(sql_upper.startswith(prefix) for prefix in allowed_prefixes): + return False, "Only SELECT queries are allowed in read-only mode" + + # Block dangerous operations + dangerous_keywords = [ + "DROP ", "DELETE ", "TRUNCATE ", "ALTER ", "CREATE ", + "INSERT ", "UPDATE ", "GRANT ", "REVOKE ", "COPY ", + "pg_", "information_schema" + ] + + if self.read_only: + for keyword in dangerous_keywords: + if keyword in sql_upper: + return False, f"Query contains forbidden keyword: {keyword.strip()}" + + return True, None + + def execute_query(self, request: AgentQueryRequest) -> AgentQueryResponse: + """ + Execute a query request from the main service. + + Args: + request: Query request with SQL and parameters + + Returns: + Query response with results or error + """ + start_time = time.time() + + try: + # Validate query + is_valid, error = self._validate_query(request.sql) + if not is_valid: + logger.warning(f"Query validation failed: {error}") + return AgentQueryResponse( + query_id=request.query_id, + success=False, + error=error + ) + + # Get connection + conn = self._get_db_connection() + + # Set schema search path + if request.schema_name and request.schema_name in self.allowed_schemas: + with conn.cursor() as cur: + cur.execute(f"SET search_path TO {request.schema_name}") + + # Execute query with timeout + with conn.cursor() as cur: + # Set statement timeout + timeout_ms = min(request.timeout_seconds, self.query_timeout) * 1000 + cur.execute(f"SET statement_timeout = {timeout_ms}") + + # Execute main query + cur.execute(request.sql, request.parameters) + + # Fetch results + max_rows = min(request.max_rows, self.max_rows) + rows = cur.fetchmany(max_rows + 1) # Fetch one extra to detect truncation + + truncated = len(rows) > max_rows + if truncated: + rows = rows[:max_rows] + + # Get column info + columns = [desc[0] for desc in cur.description] if cur.description else [] + + # Build column types + column_types = {} + if cur.description: + for desc in cur.description: + col_name = desc[0] + type_code = desc[1] + # Map PostgreSQL type codes to names + column_types[col_name] = self._get_type_name(type_code) + + execution_time = (time.time() - start_time) * 1000 + + # Update stats + self._queries_executed += 1 + self._total_execution_time += execution_time + + logger.info(f"Query {request.query_id} executed: {len(rows)} rows in {execution_time:.2f}ms") + + # Convert rows to dicts and handle non-JSON-serializable types + result_rows = [self._serialize_row(dict(row)) for row in rows] + + return AgentQueryResponse( + query_id=request.query_id, + success=True, + rows=result_rows, + row_count=len(result_rows), + columns=columns, + column_types=column_types, + execution_time_ms=execution_time, + truncated=truncated + ) + + except psycopg2.Error as e: + self._errors += 1 + execution_time = (time.time() - start_time) * 1000 + logger.error(f"Database error: {e}") + + return AgentQueryResponse( + query_id=request.query_id, + success=False, + error=f"Database error: {str(e)}", + execution_time_ms=execution_time + ) + + except Exception as e: + self._errors += 1 + execution_time = (time.time() - start_time) * 1000 + logger.error(f"Query execution error: {e}") + + return AgentQueryResponse( + query_id=request.query_id, + success=False, + error=f"Execution error: {str(e)}", + execution_time_ms=execution_time + ) + + def _get_type_name(self, type_code: int) -> str: + """Map PostgreSQL type OID to type name.""" + # Common PostgreSQL type OIDs + type_map = { + 16: "boolean", + 20: "bigint", + 21: "smallint", + 23: "integer", + 25: "text", + 700: "real", + 701: "double precision", + 1043: "varchar", + 1082: "date", + 1083: "time", + 1114: "timestamp", + 1184: "timestamptz", + 1700: "numeric", + 2950: "uuid", + 3802: "jsonb", + 114: "json", + } + return type_map.get(type_code, "unknown") + + def _serialize_row(self, row: Dict[str, Any]) -> Dict[str, Any]: + """Convert row values to JSON-serializable types.""" + from datetime import date, datetime, timedelta + from decimal import Decimal + import uuid as uuid_module + + result = {} + for key, value in row.items(): + if isinstance(value, datetime): + result[key] = value.isoformat() + elif isinstance(value, date): + result[key] = value.isoformat() + elif isinstance(value, timedelta): + result[key] = str(value) + elif isinstance(value, Decimal): + result[key] = float(value) + elif isinstance(value, uuid_module.UUID): + result[key] = str(value) + elif isinstance(value, bytes): + result[key] = value.decode('utf-8', errors='replace') + else: + result[key] = value + return result + + def verify_request(self, token: str) -> tuple[bool, Optional[str]]: + """ + Verify an incoming request token. + + Args: + token: Authentication token from request + + Returns: + Tuple of (is_valid, error_message) + """ + return validate_agent_token(token, self.agent_id, self.agent_secret) + + def get_heartbeat(self) -> AgentHeartbeat: + """Generate heartbeat message.""" + connected_dbs = [] + try: + conn = self._get_db_connection() + with conn.cursor() as cur: + cur.execute("SELECT current_database()") + result = cur.fetchone() + if result: + connected_dbs.append(result['current_database']) + except Exception: + pass + + return AgentHeartbeat( + agent_id=self.agent_id, + agent_token=self.authenticator.get_token(), + status="healthy" if connected_dbs else "degraded", + connected_databases=connected_dbs, + active_connections=1 if self._db_connection and not self._db_connection.closed else 0 + ) + + def get_status(self) -> Dict[str, Any]: + """Get agent status information.""" + uptime = None + if self._start_time: + uptime = int((datetime.utcnow() - self._start_time).total_seconds()) + + avg_time = 0.0 + if self._queries_executed > 0: + avg_time = self._total_execution_time / self._queries_executed + + error_rate = 0.0 + total = self._queries_executed + self._errors + if total > 0: + error_rate = self._errors / total + + return { + "agent_id": self.agent_id, + "version": self.VERSION, + "status": "running" if self._is_running else "stopped", + "uptime_seconds": uptime, + "queries_executed": self._queries_executed, + "avg_response_time_ms": avg_time, + "error_rate": error_rate, + "allowed_schemas": self.allowed_schemas, + "read_only": self.read_only, + "max_rows": self.max_rows + } + + def test_database_connection(self) -> tuple[bool, str]: + """ + Test the database connection. + + Returns: + Tuple of (success, message) + """ + try: + conn = self._get_db_connection() + with conn.cursor() as cur: + cur.execute("SELECT 1") + cur.fetchone() + return True, "Database connection successful" + except Exception as e: + return False, f"Database connection failed: {str(e)}" + + def start(self): + """Start the agent service.""" + self._is_running = True + self._start_time = datetime.utcnow() + logger.info(f"Agent service started: {self.agent_id}") + + def stop(self): + """Stop the agent service and close connections.""" + self._is_running = False + if self._db_connection: + self._db_connection.close() + self._db_connection = None + logger.info(f"Agent service stopped: {self.agent_id}") + + +def create_agent_from_config(config_path: str) -> AgentService: + """ + Create an agent service from a configuration file. + + Args: + config_path: Path to JSON configuration file + + Returns: + Configured AgentService instance + """ + with open(config_path, 'r') as f: + config = json.load(f) + + return AgentService( + agent_id=config["agent_id"], + agent_secret=config["agent_secret"], + database_url=config["database_url"], + allowed_schemas=config.get("allowed_schemas", ["public"]), + read_only=config.get("read_only", True), + max_rows=config.get("max_rows", 10000), + query_timeout=config.get("query_timeout", 30) + ) diff --git a/backend/main.py b/backend/main.py index 689aa9f..23de29e 100644 --- a/backend/main.py +++ b/backend/main.py @@ -22,8 +22,14 @@ from utils.datasource_loader import download_from_cloudinary, load_datasource_files, load_chat, ensure_datasource_files from utils.database_utilities import get_db_connection, db_connection, db_cursor from utils.postgres_connector import PostgresConnector +from agent.manager import init_agent_manager, get_agent_manager +from agent.models import AgentRegistration, AgentHeartbeat, AgentQueryRequest app = FastAPI() + +# Initialize agent manager +agent_manager = init_agent_manager(db_cursor) + url=os.getenv("NEXT_JS_API_URL") # Configure CORS app.add_middleware( @@ -1461,6 +1467,323 @@ def convert_for_json(obj): raise HTTPException(status_code=500, detail=str(e)) +# ============================================================================ +# CUSTOMER-HOSTED AGENT ENDPOINTS +# ============================================================================ + +@app.post("/agent/register") +async def register_agent(req: dict): + """ + Register a new customer-hosted agent. + + Request body: + { + "agent_name": "My Production Agent", + "host_url": "https://agent.mycompany.com:8443", + "database_type": "postgres", + "allowed_schemas": ["public"], + "user_id": "user_xxx" + } + + Response: + { + "agent_id": "agent_xxx", + "agent_token": "tok_xxx", + "agent_secret": "secret_xxx", + "message": "...", + "status": "pending" + } + """ + try: + user_id = req.get("user_id") + agent_name = req.get("agent_name") + host_url = req.get("host_url") + + if not all([user_id, agent_name, host_url]): + raise HTTPException( + status_code=400, + detail="user_id, agent_name, and host_url are required" + ) + + registration = AgentRegistration( + agent_name=agent_name, + host_url=host_url, + database_type=req.get("database_type", "postgres"), + allowed_schemas=req.get("allowed_schemas", ["public"]), + user_id=user_id + ) + + result = await agent_manager.register_agent(registration) + + print(f"[AGENT] ✓ Registered agent: {result.agent_id} for user: {user_id}") + + return result.model_dump() + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error registering agent: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/agent/heartbeat") +async def agent_heartbeat(req: dict): + """ + Process heartbeat from a customer-hosted agent. + + Request body: + { + "agent_id": "agent_xxx", + "agent_token": "tok_xxx", + "status": "healthy", + "connected_databases": ["production"], + "active_connections": 5 + } + + Response: + { + "success": true, + "message": "Heartbeat acknowledged" + } + """ + try: + agent_id = req.get("agent_id") + agent_token = req.get("agent_token") + + if not agent_id or not agent_token: + raise HTTPException( + status_code=400, + detail="agent_id and agent_token are required" + ) + + heartbeat = AgentHeartbeat( + agent_id=agent_id, + agent_token=agent_token, + status=req.get("status", "healthy"), + connected_databases=req.get("connected_databases", []), + active_connections=req.get("active_connections", 0) + ) + + result = await agent_manager.process_heartbeat(heartbeat) + + return result + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error processing heartbeat: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/agent/query") +async def execute_agent_query(req: dict): + """ + Execute a query through a customer-hosted agent. + + Request body: + { + "agent_id": "agent_xxx", + "user_id": "user_xxx", + "sql": "SELECT * FROM users LIMIT 10", + "parameters": {}, + "schema_name": "public", + "timeout": 30 + } + + Response: + { + "query_id": "query_xxx", + "success": true, + "rows": [...], + "row_count": 10, + "columns": ["id", "name"], + "execution_time_ms": 45.2 + } + """ + try: + agent_id = req.get("agent_id") + user_id = req.get("user_id") + sql = req.get("sql") + + if not all([agent_id, user_id, sql]): + raise HTTPException( + status_code=400, + detail="agent_id, user_id, and sql are required" + ) + + # Verify user owns the agent + agent_status = await agent_manager.get_agent_status(agent_id, user_id) + if not agent_status: + raise HTTPException(status_code=404, detail="Agent not found or unauthorized") + + result = await agent_manager.execute_query( + agent_id=agent_id, + sql=sql, + parameters=req.get("parameters"), + schema_name=req.get("schema_name", "public"), + timeout=req.get("timeout", 30) + ) + + return result.model_dump() + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error executing query: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get("/agent/status/{agent_id}") +async def get_agent_status(agent_id: str, user_id: str): + """ + Get status information for an agent. + + Query params: + - user_id: Owner user ID for authorization + + Response: + { + "agent_id": "agent_xxx", + "agent_name": "My Agent", + "status": "connected", + "host_url": "https://...", + "last_heartbeat": "2026-01-02T12:00:00Z", + "queries_executed": 100, + "avg_response_time_ms": 45.2 + } + """ + try: + if not user_id: + raise HTTPException(status_code=400, detail="user_id is required") + + result = await agent_manager.get_agent_status(agent_id, user_id) + + if not result: + raise HTTPException(status_code=404, detail="Agent not found or unauthorized") + + return result.model_dump() + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error getting agent status: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get("/agent/list") +async def list_agents(user_id: str): + """ + List all agents for a user. + + Query params: + - user_id: Owner user ID + + Response: + [ + { + "agent_id": "agent_xxx", + "agent_name": "My Agent", + "status": "connected", + ... + } + ] + """ + try: + if not user_id: + raise HTTPException(status_code=400, detail="user_id is required") + + agents = await agent_manager.list_agents(user_id) + + return [agent.model_dump() for agent in agents] + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error listing agents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete("/agent/{agent_id}") +async def delete_agent(agent_id: str, user_id: str): + """ + Delete an agent. + + Query params: + - user_id: Owner user ID for authorization + + Response: + { + "success": true, + "message": "Agent deleted successfully" + } + """ + try: + if not user_id: + raise HTTPException(status_code=400, detail="user_id is required") + + deleted = await agent_manager.delete_agent(agent_id, user_id) + + if not deleted: + raise HTTPException(status_code=404, detail="Agent not found or unauthorized") + + print(f"[AGENT] ✓ Deleted agent: {agent_id}") + + return { + "success": True, + "message": "Agent deleted successfully", + "agent_id": agent_id + } + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error deleting agent: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/agent/test-connection") +async def test_agent_connection(req: dict): + """ + Test connection to a customer-hosted agent. + + Request body: + { + "agent_id": "agent_xxx", + "user_id": "user_xxx" + } + + Response: + { + "success": true, + "latency_ms": 45.2, + "message": "Connection successful" + } + """ + try: + agent_id = req.get("agent_id") + user_id = req.get("user_id") + + if not agent_id or not user_id: + raise HTTPException( + status_code=400, + detail="agent_id and user_id are required" + ) + + success, message, latency = await agent_manager.test_agent_connection(agent_id, user_id) + + return { + "success": success, + "latency_ms": latency, + "message": message + } + + except HTTPException: + raise + except Exception as e: + print(f"[AGENT] ✗ Error testing connection: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000) \ No newline at end of file diff --git a/my-app/prisma/schema.prisma b/my-app/prisma/schema.prisma index 3713591..1b1315e 100644 --- a/my-app/prisma/schema.prisma +++ b/my-app/prisma/schema.prisma @@ -20,6 +20,7 @@ model User { dataSources DataSource[] conversations Conversation[] sessions Session[] + agents Agent[] @@index([email]) } @@ -81,3 +82,31 @@ model Session { @@index([dataSourceId]) // ✅ NEW @@index([createdAt]) } + +model Agent { + id String @id @default(cuid()) + userId String + name String + hostUrl String + databaseType String @default("postgres") + allowedSchemas Json @default("[]") + agentToken String + agentSecret String + status String @default("pending") // pending, connected, disconnected, error + lastHeartbeat DateTime? + connectedDatabases Json? @default("[]") + activeConnections Int @default(0) + queriesExecuted Int @default(0) + avgResponseTime Float @default(0) + errorRate Float @default(0) + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt + + // Relations + user User @relation(fields: [userId], references: [id], onDelete: Cascade) + + // Constraints + @@index([userId]) + @@index([status]) + @@index([createdAt]) +} From a9e56fc4c7befc329fce1f11ab4a114374b57be8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 2 Jan 2026 06:27:02 +0000 Subject: [PATCH 3/4] Fix code review issues: move imports to top, fix SQL injection vulnerability Co-authored-by: ancient-kid <183126081+ancient-kid@users.noreply.github.com> --- backend/agent/manager.py | 6 ++---- backend/agent/service.py | 16 ++++++++++------ 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/backend/agent/manager.py b/backend/agent/manager.py index 9015d8e..c413b86 100644 --- a/backend/agent/manager.py +++ b/backend/agent/manager.py @@ -7,6 +7,8 @@ import os import json +import time +import uuid import asyncio import logging import httpx @@ -175,8 +177,6 @@ async def execute_query( Returns: Query response from agent """ - import uuid - # Get agent data agent_data = await self._get_agent_data(agent_id) if not agent_data: @@ -369,8 +369,6 @@ async def test_agent_connection(self, agent_id: str, user_id: str) -> Tuple[bool Returns: Tuple of (success, message, latency_ms) """ - import time - agent_data = await self._get_agent_data(agent_id) if not agent_data: return False, "Agent not found", 0 diff --git a/backend/agent/service.py b/backend/agent/service.py index 81e7574..283972c 100644 --- a/backend/agent/service.py +++ b/backend/agent/service.py @@ -20,6 +20,7 @@ import time import asyncio import logging +import uuid as uuid_module from typing import Optional, Dict, Any, List from datetime import datetime @@ -164,16 +165,20 @@ def execute_query(self, request: AgentQueryRequest) -> AgentQueryResponse: # Get connection conn = self._get_db_connection() - # Set schema search path + # Set schema search path - validate against allowed schemas to prevent injection if request.schema_name and request.schema_name in self.allowed_schemas: + # Use psycopg2.sql module for safe identifier handling + from psycopg2 import sql with conn.cursor() as cur: - cur.execute(f"SET search_path TO {request.schema_name}") + cur.execute( + sql.SQL("SET search_path TO {}").format(sql.Identifier(request.schema_name)) + ) # Execute query with timeout with conn.cursor() as cur: - # Set statement timeout + # Set statement timeout - timeout_ms is already validated as integer timeout_ms = min(request.timeout_seconds, self.query_timeout) * 1000 - cur.execute(f"SET statement_timeout = {timeout_ms}") + cur.execute("SET statement_timeout = %s", (timeout_ms,)) # Execute main query cur.execute(request.sql, request.parameters) @@ -269,9 +274,8 @@ def _get_type_name(self, type_code: int) -> str: def _serialize_row(self, row: Dict[str, Any]) -> Dict[str, Any]: """Convert row values to JSON-serializable types.""" - from datetime import date, datetime, timedelta + from datetime import date, timedelta from decimal import Decimal - import uuid as uuid_module result = {} for key, value in row.items(): From 7f49f065f683163980b00344f880d1e638ccbd1f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 2 Jan 2026 08:43:17 +0000 Subject: [PATCH 4/4] Add Streamlit demo app with file upload and chatbot interface Co-authored-by: ancient-kid <183126081+ancient-kid@users.noreply.github.com> --- backend/demo/README.md | 79 ++++++++++ backend/demo/requirements.txt | 5 + backend/demo/streamlit_app.py | 282 ++++++++++++++++++++++++++++++++++ 3 files changed, 366 insertions(+) create mode 100644 backend/demo/README.md create mode 100644 backend/demo/requirements.txt create mode 100644 backend/demo/streamlit_app.py diff --git a/backend/demo/README.md b/backend/demo/README.md new file mode 100644 index 0000000..f5c8a50 --- /dev/null +++ b/backend/demo/README.md @@ -0,0 +1,79 @@ +# RELIX Streamlit Demo Application + +A simple demo chat interface for showcasing the NL-to-SQL capability with customer-hosted agent. + +## Features + +- 📁 Upload CSV/Excel files +- 💬 Ask natural language questions about your data +- 🔒 Secure query execution through customer-hosted agent +- 📊 View query results and insights + +## Quick Start + +### 1. Install Dependencies + +```powershell +cd backend/demo +pip install -r requirements.txt +``` + +### 2. Set Environment Variables (Optional) + +```powershell +$env:SAAS_SERVER_URL = "http://localhost:8000" +$env:AGENT_ID = "your_agent_id" +$env:USER_ID = "your_user_id" +``` + +### 3. Run the App + +```powershell +streamlit run streamlit_app.py +``` + +The app will open at `http://localhost:8501` + +## Usage + +1. **Configure Connection**: Enter your SaaS server URL and Agent ID in the sidebar +2. **Upload Data**: Upload a CSV or Excel file in the left panel +3. **Ask Questions**: Type natural language questions in the right panel +4. **View Results**: See SQL queries and results displayed in the chat + +## For Hackathon Demo + +### Two-Laptop Setup + +**Laptop 1 (SaaS Server):** +```powershell +cd backend +uvicorn main:app --host 0.0.0.0 --port 8000 +``` + +**Laptop 2 (Customer Demo):** +```powershell +cd backend/demo +$env:SAAS_SERVER_URL = "http://192.168.x.x:8000" # SaaS laptop IP +$env:AGENT_ID = "agent_xxxxx" # From registration +streamlit run streamlit_app.py +``` + +### Demo Script + +1. Show the upload feature - drag and drop a CSV file +2. Point out the schema detection in the sidebar +3. Ask a question like "Show me the top 5 rows" +4. Explain that the query goes through the secure agent, never exposing raw database credentials + +## Screenshots + +The app has two main panels: +- **Left**: File upload and data preview +- **Right**: Chat interface for questions + +## Troubleshooting + +- **Connection Error**: Make sure the backend server is running +- **Agent Not Found**: Verify agent_id is correct and agent is registered +- **No Results**: Check that your agent is connected (status: connected) diff --git a/backend/demo/requirements.txt b/backend/demo/requirements.txt new file mode 100644 index 0000000..8edc48f --- /dev/null +++ b/backend/demo/requirements.txt @@ -0,0 +1,5 @@ +# Streamlit Demo App Dependencies +streamlit>=1.28.0 +pandas>=2.0.0 +openpyxl>=3.1.0 +requests>=2.31.0 diff --git a/backend/demo/streamlit_app.py b/backend/demo/streamlit_app.py new file mode 100644 index 0000000..c6d82a2 --- /dev/null +++ b/backend/demo/streamlit_app.py @@ -0,0 +1,282 @@ +""" +RELIX Demo Application - Streamlit Chat Interface + +A simple demo application that allows users to: +1. Upload CSV/Excel files +2. Ask natural language questions about their data +3. Get SQL queries and results through the customer-hosted agent + +Usage: + pip install streamlit pandas openpyxl requests + streamlit run demo/streamlit_app.py +""" + +import streamlit as st +import pandas as pd +import requests +import json +import os + +# Configuration +SAAS_SERVER_URL = os.getenv("SAAS_SERVER_URL", "http://localhost:8000") +AGENT_ID = os.getenv("AGENT_ID", "") +USER_ID = os.getenv("USER_ID", "demo_user") + +# Page configuration +st.set_page_config( + page_title="RELIX - NL to SQL Demo", + page_icon="🔍", + layout="wide" +) + +# Custom CSS +st.markdown(""" + +""", unsafe_allow_html=True) + +# Title +st.title("🔍 RELIX - Natural Language to SQL") +st.markdown("Upload your data and ask questions in plain English!") + +# Sidebar for configuration +with st.sidebar: + st.header("⚙️ Configuration") + + server_url = st.text_input( + "SaaS Server URL", + value=SAAS_SERVER_URL, + help="URL of your RELIX SaaS server" + ) + + agent_id = st.text_input( + "Agent ID", + value=AGENT_ID, + help="Your registered agent ID" + ) + + user_id = st.text_input( + "User ID", + value=USER_ID, + help="Your user ID" + ) + + st.divider() + + st.header("📊 Agent Status") + if st.button("Check Agent Status"): + if agent_id and user_id: + try: + response = requests.get( + f"{server_url}/agent/status/{agent_id}", + params={"user_id": user_id}, + timeout=5 + ) + if response.status_code == 200: + status = response.json() + st.success(f"Status: {status.get('status', 'unknown')}") + else: + st.error(f"Error: {response.status_code}") + except Exception as e: + st.error(f"Connection error: {str(e)}") + else: + st.warning("Please enter Agent ID and User ID") + +# Initialize session state +if "messages" not in st.session_state: + st.session_state.messages = [] + +if "uploaded_data" not in st.session_state: + st.session_state.uploaded_data = None + +if "data_schema" not in st.session_state: + st.session_state.data_schema = None + +# Main content area +col1, col2 = st.columns([1, 1]) + +with col1: + st.header("📁 Upload Data") + + uploaded_file = st.file_uploader( + "Choose a CSV or Excel file", + type=["csv", "xlsx", "xls"], + help="Upload your data file to analyze" + ) + + if uploaded_file is not None: + try: + # Read the file + if uploaded_file.name.endswith('.csv'): + df = pd.read_csv(uploaded_file) + else: + df = pd.read_excel(uploaded_file) + + st.session_state.uploaded_data = df + + # Generate schema info + schema_info = [] + for col in df.columns: + dtype = str(df[col].dtype) + schema_info.append(f"- {col} ({dtype})") + st.session_state.data_schema = "\n".join(schema_info) + + st.success(f"✅ Loaded {len(df)} rows, {len(df.columns)} columns") + + # Show preview + st.subheader("Data Preview") + st.dataframe(df.head(10), use_container_width=True) + + # Show schema + with st.expander("📋 Column Schema"): + st.text(st.session_state.data_schema) + + except Exception as e: + st.error(f"Error reading file: {str(e)}") + +with col2: + st.header("💬 Ask Questions") + + # Display chat history + for message in st.session_state.messages: + if message["role"] == "user": + st.markdown(f""" +
+ You: {message["content"]} +
+ """, unsafe_allow_html=True) + else: + st.markdown(f""" +
+ Assistant:
{message["content"]} +
+ """, unsafe_allow_html=True) + + # Question input + question = st.text_area( + "Ask a question about your data", + placeholder="e.g., Show me the top 5 customers by total sales", + height=100, + key="question_input" + ) + + col_btn1, col_btn2 = st.columns([1, 1]) + + with col_btn1: + ask_button = st.button("🚀 Ask", type="primary", use_container_width=True) + + with col_btn2: + if st.button("🗑️ Clear Chat", use_container_width=True): + st.session_state.messages = [] + st.rerun() + + if ask_button and question: + # Add user message + st.session_state.messages.append({ + "role": "user", + "content": question + }) + + # Process the question + with st.spinner("Thinking..."): + try: + if agent_id and server_url: + # Build context with schema info + context = "" + if st.session_state.data_schema: + context = f"Available columns:\n{st.session_state.data_schema}\n\n" + + # Call the query endpoint + # Note: In a real implementation, this would go through the NL-to-SQL pipeline + # For demo purposes, we'll show how to structure the API call + + payload = { + "agent_id": agent_id, + "user_id": user_id, + "question": question, + "context": context + } + + # Try to call the main query endpoint + response = requests.post( + f"{server_url}/query", + json={ + "question": f"{context}User question: {question}", + "user_id": user_id, + "data_source_id": agent_id # Using agent as data source + }, + timeout=30 + ) + + if response.status_code == 200: + result = response.json() + answer = result.get("answer", result.get("final_answer", "No answer received")) + + # Format the response + response_text = f"{answer}" + + if "insights" in result and result["insights"]: + response_text += f"\n\n**Insights:** {result['insights']}" + + else: + # Fallback: Show how to use the agent/query endpoint directly + response_text = f""" +**Demo Mode**: To execute queries through your agent, use: + +``` +POST {server_url}/agent/query +{{ + "agent_id": "{agent_id}", + "user_id": "{user_id}", + "sql": "SELECT ... FROM your_table" +}} +``` + +Your question: "{question}" + +With the uploaded data schema: +{st.session_state.data_schema or "No data uploaded"} +""" + else: + response_text = "⚠️ Please configure Agent ID and Server URL in the sidebar." + + st.session_state.messages.append({ + "role": "assistant", + "content": response_text + }) + + except requests.exceptions.ConnectionError: + st.session_state.messages.append({ + "role": "assistant", + "content": f"❌ Could not connect to server at {server_url}. Make sure the backend is running." + }) + except Exception as e: + st.session_state.messages.append({ + "role": "assistant", + "content": f"❌ Error: {str(e)}" + }) + + st.rerun() + +# Footer +st.divider() +st.markdown(""" +
+

🔒 Your data stays secure with the customer-hosted agent architecture

+

Built with ❤️ using RELIX NL-to-SQL

+
+""", unsafe_allow_html=True)