import os import json import time import logging import uuid import asyncio import sys from typing import Dict, Any, List, Optional, Set from textwrap import dedent from datetime import datetime # Load environment variables from .env file from dotenv import load_dotenv load_dotenv(os.path.join(os.path.dirname(__file__), '..', '.env')) # FastAPI imports for custom tenant-aware endpoint from fastapi import FastAPI, HTTPException, Body, Depends, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel from backend.core.auth import get_current_user, AuthUser # Updated imports for comprehensive tracking from agno.db.sqlite import SqliteDb # Changed from InMemoryDb for persistence from agno.agent import Agent from agno.models.nvidia import Nvidia from agno.os import AgentOS from agno.run import RunContext from agno.run.agent import RunEvent # Import the new multi-tenant toolkit from backend.SQL_Agent.data_sources_sql_toolkit import DataSourcesSQLToolkit from backend.SQL_Agent.tenant_file_toolkit import TenantFileToolkit # Configuration for data sources API DATA_SOURCES_API_BASE_URL = os.environ.get("DATA_SOURCES_API_BASE_URL", "http://127.0.0.1:8000") DATA_SOURCES_API_KEY = os.environ.get("DATA_SOURCES_API_KEY") # Optional API key for authenticated requests print(f"šŸ“” Data Sources API URL: {DATA_SOURCES_API_BASE_URL}") if DATA_SOURCES_API_KEY: print("šŸ”‘ Data Sources API Key configured.") else: print(" No Data Sources API Key configured (optional)") logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def _get_billing_redis(): try: import redis redis_url = os.environ.get("REDIS_URL") if redis_url: return redis.from_url(redis_url, decode_responses=True) redis_host = os.environ.get("REDIS_HOST") redis_port = int(os.environ.get("REDIS_PORT", "6379")) redis_db = int(os.environ.get("REDIS_DB", "0")) redis_password = os.environ.get("REDIS_PASSWORD") if redis_host: return redis.Redis( host=redis_host, port=redis_port, db=redis_db, password=redis_password, decode_responses=True, ) return None except Exception as exc: logger.warning(f"Billing Redis unavailable: {exc}") return None def record_tenant_billing(tenant_id: str, input_tokens: int, output_tokens: int) -> None: if not tenant_id: return billing_redis = _get_billing_redis() if billing_redis is None: return input_tokens = int(input_tokens or 0) output_tokens = int(output_tokens or 0) total_tokens = input_tokens + output_tokens billing_key = f"tenant_billing:{tenant_id}" billing_redis.hincrby(billing_key, "input_tokens", input_tokens) billing_redis.hincrby(billing_key, "output_tokens", output_tokens) billing_redis.hincrby(billing_key, "total_tokens", total_tokens) est_cost = (input_tokens / 1_000_000) * 0.15 + (output_tokens / 1_000_000) * 0.60 billing_redis.hincrbyfloat(billing_key, "estimated_cost_usd", float(f"{est_cost:.6f}")) # NEW: Enhanced Tool Hook for Complete Logging def comprehensive_logging_hook( run_context: RunContext, function_name: str, function_call, arguments: Dict[str, Any] ) -> Any: """ Comprehensive tool execution logging hook that saves: - Tool name and arguments - Execution timestamp - Results - User context """ # Access session_state from run_context (Agno v2 API) if not run_context.session_state: run_context.session_state = {} session_state = run_context.session_state # Initialize logging structure in session state if "tool_execution_log" not in session_state: session_state["tool_execution_log"] = [] # Create execution record execution_start = datetime.now() execution_record = { "tool_name": function_name, "arguments": arguments, "timestamp": execution_start.isoformat(), "execution_id": f"{function_name}_{execution_start.timestamp()}" } logger.info(f"šŸ”§ Executing tool: {function_name} with args: {arguments}") try: # Execute the actual tool result = function_call(**arguments) # Log successful execution execution_end = datetime.now() execution_record.update({ "result": str(result)[:1000], # Truncate long results "status": "success", "duration_ms": (execution_end - execution_start).total_seconds() * 1000, "completed_at": execution_end.isoformat() }) logger.info(f"āœ… Tool {function_name} completed successfully in {execution_record['duration_ms']:.2f}ms") except Exception as e: # Log failed execution execution_end = datetime.now() execution_record.update({ "error": str(e), "status": "failed", "duration_ms": (execution_end - execution_start).total_seconds() * 1000, "completed_at": execution_end.isoformat() }) logger.error(f"āŒ Tool {function_name} failed: {str(e)}") raise # Re-raise the exception finally: # Always save the execution record session_state["tool_execution_log"].append(execution_record) return result system_prompt = dedent(""" Sirus PhobosQ Sirus The Data Scientist & Strategist Bridge the gap between raw database rows and high-level business strategy. Professional, energetic, precise, and helpful. You speak in Markdown. The user CANNOT see your tool calls, JSON outputs, or SQL code. You MUST translate every tool result into a natural language sentence. NEVER end a turn with a tool call. ALWAYS end with a text response. Your semantic search is strict. When searching for tables, you MUST expand keywords. - If user asks: "How many users?" -> Search: ['users', 'accounts', 'customers', 'profiles','people','members'etc...] - If user asks: "Sales?" -> Search: ['sales', 'orders', 'transactions', 'revenue', 'invoices','bookings'] If `find_relevant_tables` returns 0 matches, you MUST NOT give up. You MUST immediately call `get_available_sources_and_schema` to pull the full database map. Then, manually find the table and execute the query. NEVER execute INSERT, UPDATE, DELETE, DROP, or ALTER. ALWAYS use `LIMIT 100` on list queries to prevent token overflows. Do I have the `source_instructions` in my context? If NO: Call `list_sources`, select the most relevant one, then `get_source_instructions`. If YES: Skip to Phase 2. Call `find_relevant_tables(question, concepts)`. Use broad concepts. If the user asks a "Why" question, search for fact tables (orders, logs) AND dimension tables (users, products). If matches == 0: Call `get_available_sources_and_schema(tenant_id)`. Call `execute_sql_query(sql_query)` OR `save_query_to_tenant_csv(sql_query)`. 1. Write Standard ANSI SQL. 2. Use the exact table names found in Phase 2. 3. If the user asks "Why" or "Trend", run aggregations (GROUP BY). 4. **CRITICAL ML RULE:** If the user asks to "save", "export", "analyze in pandas", or "prepare for ML", you MUST use `save_query_to_tenant_csv`. If SQL fails: Read error -> Correct Syntax -> Retry Query. Convert JSON list to Text. Start of conversation or when switching databases. Finds the tenant_id and source_name. Immediately after picking a source. Gets the "Manual" for the database (SQL dialect, special column rules). Every user question. Argument `concepts` must be a list of broad synonyms. Example: User="Churn rate?" -> concepts=["churn", "status", "active", "cancelled", "users"] ONLY when `find_relevant_tables` fails (returns []). The "Nuclear Option". Dumps the whole schema so you can find tables manually. Once you have table names and a clear intent for a simple data pull or counting. - No Markdown in the SQL string. - Dates should be handled dynamically (e.g., `CURRENT_DATE`). - Always handle NULLs in math operations (`COALESCE`). When a user asks to export data, prepare it for Machine Learning, or save it. Executes SQL but completely bypasses standard memory limits by saving directly to the MinIO cluster. When a user asks what files, reports, or datasets they have in their workspace. When a user asks you to analyze a specific CSV file in their workspace. How many users are there? 1. **Init**: I have instructions. 2. **Search**: `find_relevant_tables` for "users", "accounts". 3. **Result**: Found table `public.users`. 4. **Execute**: `SELECT COUNT(*) as count FROM public.users`. 5. **Result**: `[{'count': 5420}]`. 6. **Speak**: "There are currently **5,420 users** in your database. Is there a specific segment you'd like to analyze?" Why did sales drop last month? 1. **Search**: `find_relevant_tables` for "sales", "orders", "revenue". 2. **Result**: Found `orders`. 3. **Reasoning**: To find "why", I need to look at categories or regions over time. A simple sum isn't enough. 4. **Execute**: `SELECT category, SUM(amount) as total FROM orders WHERE date >= DATE('now', '-1 month') GROUP BY category ORDER BY total ASC`. 5. **Result**: `[{'category': 'Tech', 'total': 500}, {'category': 'Home', 'total': 50000}]`. 6. **Speak**: "I analyzed the sales drop. It appears **Tech** products severely underperformed ($500) compared to Home Goods. **Strategic Insight:** This often indicates a stockout or a lack of promotion in that category. **Recommendation:** I suggest checking inventory levels for Tech products immediately." Show me the invoices. 1. **Search**: `find_relevant_tables(concepts=['invoices', 'billing'])`. 2. **Result**: `Matches: []` (Zero tables found). 3. **Pivot**: The vector search failed. I must see the schema. 4. **Tool**: `get_available_sources_and_schema(tenant_id=...)`. 5. **Result**: Full Schema JSON. I read it. I see a table named `billing_ledgers`. 6. **Execute**: `SELECT * FROM billing_ledgers LIMIT 5`. 7. **Speak**: "I couldn't find a table explicitly named 'invoices', but I found `billing_ledgers` which contains billing data. Here are the top 5 records..." Analyze my demographic data file. 1. **Search**: I don't know the exact file name. I need to list them. 2. **Tool**: `list_tenant_assets()`. 3. **Result**: `[{"asset_id": "123", "filename": "Demographic Data.xlsx"}]`. 4. **Tool**: `load_tenant_file_to_dataframe(asset_id="123", instructions="Summarize the key demographics.")`. 5. **Result**: Pandas output showing age distributions and counts. 6. **Speak**: "I've analyzed your demographic dataset. The majority of the entries fall into..." Compare the regional sales in the DB with the targets in the targets.xlsx file. 1. **DB Step**: I need regional sales from the database. 2. **Tool**: `find_relevant_tables(concepts=['sales', 'regions', 'geography'])`. 3. **Execute**: `SELECT region, SUM(amount) FROM regional_sales GROUP BY region`. 4. **File Step**: Now I need the targets. 5. **Tool**: `list_tenant_assets()`. 6. **Tool**: `load_tenant_file_to_dataframe(filename="targets.xlsx", instructions="Extract the targets by region.")`. 7. **Synthesize**: I combine the SQL results and Pandas results in my reasoning. 8. **Speak**: "Comparing your database records to the Excel file, the North region hit its target ($50k vs $45k expected), but..." What does the file users.csv say? 1. **BAD PATH**: Calling `execute_sql_query('SELECT * FROM "users.csv"')`. (This is WRONG! It's a file, not a table). 2. **CORRECT PATH**: The user specifically said "file" and "csv". 3. **Tool**: `list_tenant_assets()` to verify it exists. 4. **Tool**: `load_tenant_file_to_dataframe(filename="users.csv", ...)`. 5. **Speak**: "I loaded the users.csv file and found 300 entries..." Can you show me some insights? 1. **Recognize Ambiguity**: The user never said "table" or "file". I must sample both data sources. 2. **DB Probe**: `find_relevant_tables(concepts=['users','sales','orders','activity','logs'])` to surface likely tables. 3. **File Probe**: `list_tenant_assets()` to see if any CSV/Excel assets exist that look relevant (recent uploads, names with "report", "data", etc.). 4. **Pick One of Each (Lightweight)**: Grab a tiny preview: `execute_sql_query('SELECT * FROM LIMIT 5')` and `load_tenant_file_to_dataframe(asset_id=, instructions="Give me a quick summary")`. 5. **Synthesize**: Combine the quick peeks and present the clearest starting point. Offer options: continue with DB analysis, or dive into the file. 6. **Speak**: "I checked both your database and uploaded files. From the database, I saw a table with recent activity; from files, there's a recent report.xlsx. Which one should I dig into further?" When a user explicitly mentions a "file", "csv", "excel", "dataset", or "xlsx", you MUST prioritize file-based tools (`list_tenant_assets`, `load_tenant_file_to_dataframe`). DO NOT try to query these files using standard SQL tools unless explicitly attached as a temporary table (which they are not). Files live in a separate blob storage; databases live in SQL. If the user asks a question that spans both their database AND an uploaded file: 1. Extract the DB information first using Phase 2 (Discovery) and Phase 3 (Execution). 2. Extract the File information second by locating the file with `list_tenant_assets` and querying it with `load_tenant_file_to_dataframe`. 3. Synthesize the findings using your own reasoning to combine the disparate data sources. If a tool call to `load_tenant_file_to_dataframe` fails with "File not found" or "NoSuchKey", DO NOT confidently report that the data is missing. Instead, ALWAYS call `list_tenant_assets` to double-check the exact spelling, path, or `asset_id` of the available files and try again using the exact identifier. if u encounter any errors , kindly rectify them and proceed with the task at hand. if still its an server error or something , just say that kindly neatly. - Use **Bold** for numbers and key entities. - Use Tables for lists of data. - Be concise but friendly. - CRITICAL: You MUST ALWAYS append exactly 3 follow-up questions under the exact markdown heading `### Suggested Questions` at the very end of your response. """) print("āœ… Configuration set. Initializing enhanced agent with comprehensive logging...") # Define agent IDs for AgentOS DEFAULT_AGENT_OS_ID = os.getenv("SQL_AGENT_OS_ID", "sql-agent-os") DEFAULT_AGENT_ID = os.getenv("SQL_AGENT_ID", "sirus-sql-agent") IS_PYTEST = "PYTEST_CURRENT_TEST" in os.environ or "pytest" in sys.modules agent_db = None data_sources_sql_toolkit = None tenant_file_toolkit = None gemini_sql_agent = None if not IS_PYTEST: agent_db = SqliteDb(db_file="agent_sessions.db") data_sources_sql_toolkit = DataSourcesSQLToolkit( api_base_url=DATA_SOURCES_API_BASE_URL, api_key=DATA_SOURCES_API_KEY ) tenant_file_toolkit = TenantFileToolkit() gemini_sql_agent = Agent( model=Nvidia( id="stepfun-ai/step-3.7-flash", #id="nvidia/nemotron-3-super-120b-a12b", max_tokens=32768, temperature=0.2, top_p=0.95 ), instructions=system_prompt, tools=[data_sources_sql_toolkit, tenant_file_toolkit], tool_hooks=[comprehensive_logging_hook], tool_call_limit=100, debug_mode=True, telemetry=False, db=agent_db, add_history_to_context=True, num_history_runs=3, read_chat_history=True, session_state={ "tool_execution_log": [], "user_context": {}, "analysis_metadata": {} }, add_session_state_to_context=True, markdown=True, add_datetime_to_context=True, exponential_backoff=True, delay_between_retries=10 ) gemini_sql_agent.id = DEFAULT_AGENT_ID data_sources_sql_toolkit.set_agent_ref(gemini_sql_agent) logger.info("Agent reference set in toolkit - session_state injection enabled") else: logger.info("Running under pytest: skipping heavy SQL agent runtime initialization") # Define Pydantic model for tenant-aware API requests class TenantRunRequest(BaseModel): """ Request model for our custom tenant-aware endpoint. This ensures all tenant context is provided in a single, secure request. Supports multi-source agent auto-detection when available_sources is provided. """ message: str supabase_jwt: str # JWT token for auth tenant_id: str # Extracted from JWT claims source_name: str # Default/primary source for query execution session_id: Optional[str] = None user_id: Optional[str] = None available_sources: Optional[list] = None # All available sources for agent auto-detection stream: bool = False background: bool = False # Define the tenant-aware endpoint function (will be added to AgentOS app later) async def run_tenant_agent( agent_id: str, request: TenantRunRequest, auth_user: AuthUser = Depends(get_current_user), http_request: Request = None, ): """ Custom endpoint to run an agent with tenant_id, source_name, and supabase_jwt injected directly into the session_state. This is the PRIMARY endpoint for multi-tenant agent execution. It ensures proper tenant isolation and security by: 1. Accepting all tenant context in the request body 2. Injecting it into session_state (not shared between requests) 3. Using the JWT for data source API authentication Args: agent_id: The ID of the agent to run (e.g., "sirus-sql-agent") request: TenantRunRequest containing all tenant context Returns: StreamingResponse (if stream=True) or direct JSON response """ # Get agent from the global agent we created agent = gemini_sql_agent if agent_id == DEFAULT_AGENT_ID else None if not agent: raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found.") # Resolve tenant/user context from validated JWT claims. resolved_tenant_id = (auth_user.tenant_id or "").strip() if not resolved_tenant_id: raise HTTPException(status_code=401, detail="Missing tenant_id in JWT claims") if request.tenant_id and request.tenant_id != resolved_tenant_id: logger.warning( f"Rejecting tenant-run due to tenant mismatch. body={request.tenant_id} jwt={resolved_tenant_id}" ) raise HTTPException(status_code=403, detail="tenant_id mismatch with authenticated user") resolved_actor_user_id = auth_user.id or request.user_id resolved_session_owner_id = resolved_tenant_id resolved_jwt = request.supabase_jwt if http_request is not None: auth_header = http_request.headers.get("Authorization", "") if auth_header.startswith("Bearer "): resolved_jwt = auth_header.replace("Bearer ", "", 1).strip() or resolved_jwt # CRITICAL: This is the state that will be loaded *for this run only*. # This is the correct, request-safe way to handle per-run context. # Each request gets its own isolated session_state. initial_state = { "supabase_jwt": resolved_jwt, # JWT for backend API auth "tenant_id": resolved_tenant_id, # Tenant context for toolkit "source_name": request.source_name, "user_id": resolved_session_owner_id, "actor_user_id": resolved_actor_user_id, "available_sources": request.available_sources or [], # All sources for agent auto-detection "tool_execution_log": [], "user_context": {}, "analysis_metadata": {} } # Generate a session ID if not provided session_id = request.session_id or str(uuid.uuid4()) logger.info(f"šŸš€ Starting tenant run for tenant_id={resolved_tenant_id}, source={request.source_name}, session={session_id}") if request.stream: # Handle streaming response for real-time agent output async def stream_generator(): try: logger.info(f"šŸŽ¬ Starting streaming for session {session_id}, message: {request.message[:50]}...") # Emit a canonical start event so frontend can persist session_id and retain memory across turns. run_started_payload = { "event": "RunStarted", "session_id": session_id, "agent_id": agent_id, } yield f"event: RunStarted\ndata: {json.dumps(run_started_payload)}\n\n" logger.info(f" āœ… Yielded RunStarted with session_id={session_id}") # agent.run returns a generator in stream mode response_generator = agent.run( request.message, stream=True, stream_events=True, # Enable full event streaming for tool calls session_id=session_id, user_id=resolved_session_owner_id, # Tag session with owner so get_session(user_id=) works session_state=initial_state ) chunk_count = 0 for chunk in response_generator: chunk_count += 1 # Handle RunEvent types for proper streaming if hasattr(chunk, 'event'): logger.info(f" [Chunk {chunk_count}] Event: {chunk.event}") if chunk.event == RunEvent.run_content: # Model text response event_data = {"content": chunk.content if hasattr(chunk, 'content') else str(chunk)} sse_event = f"event: RunContent\ndata: {json.dumps(event_data)}\n\n" yield sse_event logger.info(f" āœ… Yielded RunContent event") elif chunk.event == RunEvent.tool_call_started: # Tool starting - send full tool object for frontend tool_obj = chunk.tool if hasattr(chunk, 'tool') else None tool_call_id = getattr(tool_obj, 'tool_call_id', None) or str(uuid.uuid4()) tool_name = getattr(tool_obj, 'tool_name', 'unknown') tool_args = {} if hasattr(tool_obj, 'tool_args') and tool_obj.tool_args: tool_args = tool_obj.tool_args if isinstance(tool_obj.tool_args, dict) else {} event_data = { "tool": { "tool_call_id": tool_call_id, "tool_name": tool_name, "tool_args": tool_args, "role": "tool", "tool_call_error": False, "content": None, "metrics": {"time": 0}, "created_at": int(time.time()) }, "status": "started" } sse_event = f"event: ToolCallStarted\ndata: {json.dumps(event_data, default=str)}\n\n" yield sse_event logger.info(f" āœ… Yielded ToolCallStarted: {tool_name} (id: {tool_call_id})") elif chunk.event == RunEvent.tool_call_completed: # Tool finished - send full tool object with result tool_obj = chunk.tool if hasattr(chunk, 'tool') else None tool_call_id = getattr(tool_obj, 'tool_call_id', None) or str(uuid.uuid4()) tool_name = getattr(tool_obj, 'tool_name', 'unknown') tool_args = {} if hasattr(tool_obj, 'tool_args') and tool_obj.tool_args: tool_args = tool_obj.tool_args if isinstance(tool_obj.tool_args, dict) else {} # Get the ACTUAL tool result as a raw object (not pre-serialized) # This ensures proper JSON encoding when we serialize event_data content = None content_source = "none" if tool_obj: # Try to get the actual result from tool_obj if hasattr(tool_obj, 'result') and tool_obj.result is not None: result = tool_obj.result content_source = "tool_obj.result" # Keep as raw object for proper serialization if isinstance(result, (dict, list)): content = result # Raw object - will be serialized by outer json.dumps elif isinstance(result, str): # Try to parse if it's already JSON try: content = json.loads(result) except: content = result # Keep as string else: content = str(result) elif hasattr(tool_obj, 'content') and tool_obj.content is not None: tc = tool_obj.content content_source = "tool_obj.content" if isinstance(tc, (dict, list)): content = tc elif isinstance(tc, str): try: content = json.loads(tc) except: content = tc else: content = str(tc) # Last fallback - use chunk.content (formatted message) if content is None and hasattr(chunk, 'content') and chunk.content: content = str(chunk.content)[:2000] content_source = "chunk.content" tool_error = getattr(tool_obj, 'tool_call_error', False) if tool_obj else False exec_time = getattr(tool_obj, 'metrics', {}) if hasattr(exec_time, 'time'): exec_time = exec_time.time elif isinstance(exec_time, dict): exec_time = exec_time.get('time', 0) else: exec_time = 0 event_data = { "tool": { "tool_call_id": tool_call_id, "tool_name": tool_name, "tool_args": tool_args, "role": "tool", "tool_call_error": tool_error, "content": content, # Raw object - properly serialized by json.dumps below "metrics": {"time": exec_time}, "created_at": int(time.time()) }, "status": "completed" } sse_event = f"event: ToolCallCompleted\ndata: {json.dumps(event_data, default=str)}\n\n" yield sse_event logger.info(f" āœ… Yielded ToolCallCompleted: {tool_name} (id: {tool_call_id}) content_source: {content_source} content_type: {type(content).__name__}") elif chunk.event == RunEvent.run_completed: # Run completed - send metrics metrics_data = {} if hasattr(chunk, 'metrics') and chunk.metrics: # MiniMax M2.5 Estimated Pricing (e.g., $0.15/1M in, $0.60/1M out) input_tokens = getattr(chunk.metrics, 'input_tokens', 0) output_tokens = getattr(chunk.metrics, 'output_tokens', 0) total_tokens = getattr(chunk.metrics, 'total_tokens', 0) est_cost = (input_tokens / 1_000_000) * 0.15 + (output_tokens / 1_000_000) * 0.60 metrics_data = { "input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": total_tokens, "time_to_first_token": getattr(chunk.metrics, 'time_to_first_token', 0), "tokens_per_second": getattr(chunk.metrics, 'tokens_per_second', 0), "estimated_cost_usd": float(f"{est_cost:.6f}") } # Log the comprehensive cost tracking logger.info(f"šŸ’° [COST TRACKING] Session: {session_id} | Tenant: {resolved_tenant_id} | " f"Tokens: {input_tokens} In, {output_tokens} Out, {total_tokens} Total | " f"Est. Cost: ${est_cost:.6f}") record_tenant_billing(resolved_tenant_id, input_tokens, output_tokens) # Save to session state analysis metadata initial_state["analysis_metadata"]["final_cost_metrics"] = metrics_data event_data = {"metrics": metrics_data, "session_id": session_id} sse_event = f"event: RunCompleted\ndata: {json.dumps(event_data)}\n\n" yield sse_event logger.info(f" āœ… Yielded RunCompleted with metrics: {metrics_data}") else: # Other event types logger.info(f" āš ļø Unhandled event type: {chunk.event}") await asyncio.sleep(0.001) continue # Fallback for dict-based chunks if isinstance(chunk, dict): event = chunk.get("event") data = chunk.get("data") if event: sse_event = f"event: {event}\ndata: {json.dumps(data)}\n\n" else: sse_event = f"data: {json.dumps(chunk)}\n\n" yield sse_event logger.info(f" āœ… Yielded event: {event or 'data-only'}") # Small delay to ensure chunk is flushed before next one await asyncio.sleep(0.001) else: # Handle Pydantic objects or other objects try: logger.info(f"Processing chunk type: {type(chunk)}") # Try multiple serialization methods chunk_dict = None # Method 1: Pydantic v2 model_dump() if hasattr(chunk, 'model_dump'): try: chunk_dict = chunk.model_dump() logger.info(f"āœ… Serialized with model_dump()") except Exception as e: logger.info(f"model_dump() failed: {e}") # Method 2: Pydantic v1 dict() if chunk_dict is None and hasattr(chunk, 'dict'): try: chunk_dict = chunk.dict() logger.info(f"āœ… Serialized with dict()") except Exception as e: logger.info(f"dict() failed: {e}") # Method 3: Check if it's a Pydantic BaseModel if chunk_dict is None: try: # Try to import and check from pydantic import BaseModel if isinstance(chunk, BaseModel): chunk_dict = chunk.model_dump() logger.info(f"āœ… Serialized BaseModel with model_dump()") except Exception as e: logger.info(f"BaseModel check failed: {e}") # Method 4: Fall back to __dict__ if chunk_dict is None and hasattr(chunk, '__dict__'): chunk_dict = chunk.__dict__ logger.info(f"āœ… Serialized with __dict__") # Method 5: Last resort - convert to string if chunk_dict is None: logger.warning(f"Could not serialize chunk, converting to string: {type(chunk)}") chunk_dict = {"content": str(chunk)} # Extract event type if present event_type = chunk_dict.get("event") if event_type: logger.info(f"Sending event: {event_type}") # Debug: Show content for ReasoningStep events if event_type == "ReasoningStep": logger.info(f" ReasoningStep content: reasoning={chunk_dict.get('reasoning')}, content={chunk_dict.get('content')}, result={chunk_dict.get('result')}") logger.info(f" Full ReasoningStep dict keys: {list(chunk_dict.keys())}") # Use custom serializer that properly handles nested objects def serialize_value(obj): """Recursively serialize objects, converting to strings only when necessary""" if isinstance(obj, dict): return {k: serialize_value(v) for k, v in obj.items()} elif isinstance(obj, (list, tuple)): return [serialize_value(v) for v in obj] elif hasattr(obj, 'model_dump'): return serialize_value(obj.model_dump()) elif hasattr(obj, '__dict__') and not isinstance(obj, (str, int, float, bool, type(None))): return serialize_value(obj.__dict__) else: return obj serialized_dict = serialize_value(chunk_dict) # Special handling for ReasoningStep: convert content object to string if event_type == "ReasoningStep" and isinstance(serialized_dict.get("content"), dict): # Content is a reasoning object - serialize it as string for frontend reasoning_obj = serialized_dict.pop("content") serialized_dict["reasoning_content"] = json.dumps(reasoning_obj, default=str, ensure_ascii=False) logger.info(f" āœ… Converted ReasoningStep content to reasoning_content string") sse_event = f"event: {event_type}\ndata: {json.dumps(serialized_dict, default=str, ensure_ascii=False)}\n\n" else: logger.info(f"Sending data without event type") def serialize_value(obj): """Recursively serialize objects, converting to strings only when necessary""" if isinstance(obj, dict): return {k: serialize_value(v) for k, v in obj.items()} elif isinstance(obj, (list, tuple)): return [serialize_value(v) for v in obj] elif hasattr(obj, 'model_dump'): return serialize_value(obj.model_dump()) elif hasattr(obj, '__dict__') and not isinstance(obj, (str, int, float, bool, type(None))): return serialize_value(obj.__dict__) else: return obj serialized_dict = serialize_value(chunk_dict) sse_event = f"data: {json.dumps(serialized_dict, default=str, ensure_ascii=False)}\n\n" yield sse_event logger.info(f" āœ… Yielded event: {event_type or 'data-only'}") # Small delay to ensure chunk is flushed before next one await asyncio.sleep(0.001) except Exception as e: logger.error(f"Failed to serialize chunk: {e}, chunk type: {type(chunk)}", exc_info=True) yield f"data: {json.dumps({'error': str(e), 'content': str(chunk)}, default=str)}\n\n" await asyncio.sleep(0.001) logger.info(f"āœ… Streaming run completed for session {session_id} - sent {chunk_count} chunks") except Exception as e: logger.error(f"āŒ Error during stream generation for session {session_id}: {e}", exc_info=True) error_data = {"error": str(e), "code": "STREAM_ERROR"} yield f"event: error\ndata: {json.dumps(error_data)}\n\n" return StreamingResponse(stream_generator(), media_type="text/event-stream") else: # Handle non-streaming (blocking) response try: response = agent.run( request.message, stream=False, session_id=session_id, user_id=resolved_session_owner_id, session_state=initial_state ) # Non-streaming Cost tracking metrics_data = {} if hasattr(response, 'metrics') and response.metrics: input_tokens = getattr(response.metrics, 'input_tokens', 0) output_tokens = getattr(response.metrics, 'output_tokens', 0) total_tokens = getattr(response.metrics, 'total_tokens', 0) est_cost = (input_tokens / 1_000_000) * 0.15 + (output_tokens / 1_000_000) * 0.60 logger.info(f"šŸ’° [COST TRACKING] Session: {session_id} | Tenant: {resolved_tenant_id} | " f"Tokens: {input_tokens} In, {output_tokens} Out, {total_tokens} Total | " f"Est. Cost: ${est_cost:.6f}") record_tenant_billing(resolved_tenant_id, input_tokens, output_tokens) logger.info(f"āœ… Non-streaming run completed for session {session_id}") # The final response from agent.run is the message content return { "session_id": session_id, "tenant_id": resolved_tenant_id, "response": response } except Exception as e: logger.error(f"āŒ Error during non-streaming agent run for session {session_id}: {e}") raise HTTPException(status_code=500, detail=str(e)) if gemini_sql_agent is not None: agent_os = AgentOS( agents=[gemini_sql_agent], description="Multi-tenant SQL Agent for querying data sources across tenants." ) agentOS_app = agent_os.get_app() agentOS_app.add_api_route( "/tenant-run/{agent_id}", run_tenant_agent, methods=["POST"], name="run_tenant_agent" ) app = agentOS_app else: agent_os = None app = FastAPI() # ============================================================================ # Chat CRUD Endpoints (from agentOS_crud.md) # ============================================================================ def _serialize_session_obj(session_obj: Any) -> Dict[str, Any]: if hasattr(session_obj, "model_dump"): data = session_obj.model_dump() if isinstance(data, dict): return data if isinstance(session_obj, dict): return session_obj if hasattr(session_obj, "__dict__"): return dict(session_obj.__dict__) return {"value": str(session_obj)} def _extract_session_id(session_payload: Dict[str, Any]) -> str: return str( session_payload.get("session_id") or session_payload.get("id") or session_payload.get("sessionId") or "" ) def _ensure_agent_runtime_ready() -> None: if agent_db is None or gemini_sql_agent is None: raise HTTPException(status_code=503, detail="Agent runtime is not initialized") def _session_belongs_to_user(session_id: str, user_id: str) -> bool: """Check that session_id belongs to user_id. Two-pass approach for robustness: 1. Try get_session(user_id=user_id) ā€” works for sessions saved with user_id. 2. Fall back to get_session() without user_id and verify the stored user_id matches (or is unset, which we allow for legacy sessions). """ if not session_id or not user_id or gemini_sql_agent is None: return False try: # Pass 1: user-scoped lookup (ideal path) session_obj = gemini_sql_agent.get_session(session_id=session_id, user_id=user_id) if session_obj is not None: payload = _serialize_session_obj(session_obj) resolved_session_id = _extract_session_id(payload) return bool(resolved_session_id and resolved_session_id == session_id) # Pass 2: session-only lookup ā€” handles sessions where user_id was not saved session_obj = gemini_sql_agent.get_session(session_id=session_id) if session_obj is None: return False payload = _serialize_session_obj(session_obj) resolved_session_id = _extract_session_id(payload) if not resolved_session_id or resolved_session_id != session_id: return False # Accept if stored user_id matches OR is blank (legacy / first run before fix) stored_uid = str(payload.get("user_id") or "").strip() return (not stored_uid) or (stored_uid == user_id) except Exception as exc: logger.error(f"Failed ownership check for session {session_id}: {exc}") return False def _serialize_chat_message_sql(m) -> Dict[str, Any]: """Serialize an Agno Message to a rich dict for frontend turn reconstruction. Returns role, content, tool_calls (LLM call requests on assistant msgs), tool_call_id / tool_name (on tool-result msgs), and created_at. The frontend uses these to rebuild the streaming-equivalent ChatMessage structure (tool_calls + sqlExecutions) from DB history. """ role = str(getattr(m, "role", "user") or "user") raw_content = getattr(m, "content", None) if raw_content is None: content = "" elif isinstance(raw_content, str): content = raw_content else: try: content = json.dumps(raw_content) except Exception: content = str(raw_content) result: Dict[str, Any] = {"role": role, "content": content} created_at = getattr(m, "created_at", None) if created_at is not None: result["created_at"] = created_at # tool_calls: present on assistant messages that requested tool calls tool_calls = getattr(m, "tool_calls", None) if tool_calls: serialized_tcs = [] for tc in tool_calls: try: if isinstance(tc, dict): tc_id = str(tc.get("id") or "") fn = tc.get("function") or {} fn_name = str(fn.get("name") or "") fn_args = str(fn.get("arguments") or "{}") tc_type = str(tc.get("type") or "function") else: tc_id = str(getattr(tc, "id", "") or "") fn_obj = getattr(tc, "function", None) fn_name = str(getattr(fn_obj, "name", "") if fn_obj else "") fn_args = str(getattr(fn_obj, "arguments", "{}") if fn_obj else "{}") tc_type = str(getattr(tc, "type", "function") or "function") serialized_tcs.append({"id": tc_id, "type": tc_type, "function": {"name": fn_name, "arguments": fn_args}}) except Exception: continue if serialized_tcs: result["tool_calls"] = serialized_tcs # tool_call_id + tool_name: on tool-role messages (the result) tool_call_id = getattr(m, "tool_call_id", None) if tool_call_id: result["tool_call_id"] = str(tool_call_id) name = getattr(m, "name", None) if name: result["tool_name"] = str(name) return result @app.get("/chats/{user_id}") async def list_user_sessions(user_id: str, auth_user: AuthUser = Depends(get_current_user)): """LIST sessions for a user.""" _ensure_agent_runtime_ready() requester_tenant_id = (auth_user.tenant_id or "").strip() if not requester_tenant_id: raise HTTPException(status_code=401, detail="Missing tenant_id in JWT claims") if requester_tenant_id != user_id: raise HTTPException(status_code=403, detail="Forbidden: tenant_id mismatch") try: sessions = agent_db.get_sessions(user_id=user_id, component_id=DEFAULT_AGENT_ID, limit=200) serialized = [_serialize_session_obj(s) for s in (sessions or [])] # Enrich each session with normalised fields the frontend sidebar needs enriched = [] for s in serialized: sid = _extract_session_id(s) enriched.append({ **s, "session_id": sid, "name": s.get("session_name") or s.get("name") or f"Chat {sid[:8]}", "created_at": s.get("created_at"), }) return {"sessions": enriched} except Exception as e: logger.error(f"Failed to list sessions for user {user_id}: {e}") return {"sessions": []} @app.get("/chats/{user_id}/{session_id}") async def get_chat(user_id: str, session_id: str, auth_user: AuthUser = Depends(get_current_user)): """GET chat history for a session ā€” returns rich message data including tool call info.""" _ensure_agent_runtime_ready() requester_tenant_id = (auth_user.tenant_id or "").strip() if not requester_tenant_id: raise HTTPException(status_code=401, detail="Missing tenant_id in JWT claims") if requester_tenant_id != user_id: raise HTTPException(status_code=403, detail="Forbidden: tenant_id mismatch") if not _session_belongs_to_user(session_id=session_id, user_id=user_id): raise HTTPException(status_code=404, detail="Chat session not found") try: chat = gemini_sql_agent.get_chat_history(session_id=session_id) if not chat: return {"messages": [], "status": "completed"} return { "messages": [_serialize_chat_message_sql(m) for m in chat], "status": "completed", } except Exception as e: logger.error(f"Failed to get chat for session {session_id}: {e}") raise HTTPException(status_code=500, detail="Failed to retrieve chat history") @app.delete("/chats/{session_id}") async def delete_chat(session_id: str, auth_user: AuthUser = Depends(get_current_user)): """DELETE a session (and all its runs).""" _ensure_agent_runtime_ready() requester_tenant_id = (auth_user.tenant_id or "").strip() if not requester_tenant_id: raise HTTPException(status_code=401, detail="Missing tenant_id in JWT claims") if not _session_belongs_to_user(session_id=session_id, user_id=requester_tenant_id): raise HTTPException(status_code=404, detail="Chat session not found") try: gemini_sql_agent.delete_session(session_id=session_id, user_id=requester_tenant_id) return {"status": "deleted"} except Exception as e: logger.error(f"Failed to delete session {session_id}: {e}") raise HTTPException(status_code=500, detail="Failed to delete session") @app.post("/chats/{session_id}/rename") async def rename_chat(session_id: str, name: str = Body(..., embed=True), auth_user: AuthUser = Depends(get_current_user)): """RENAME a session.""" _ensure_agent_runtime_ready() requester_tenant_id = (auth_user.tenant_id or "").strip() if not requester_tenant_id: raise HTTPException(status_code=401, detail="Missing tenant_id in JWT claims") if not _session_belongs_to_user(session_id=session_id, user_id=requester_tenant_id): raise HTTPException(status_code=404, detail="Chat session not found") try: gemini_sql_agent.set_session_name(session_id=session_id, session_name=name) return {"status": "renamed"} except Exception as e: logger.error(f"Failed to rename session {session_id}: {e}") raise HTTPException(status_code=500, detail="Failed to rename session") @app.post("/chats/{session_id}/cancel/{run_id}") async def cancel_run(session_id: str, run_id: str, auth_user: AuthUser = Depends(get_current_user)): """CANCEL a running agent.""" _ensure_agent_runtime_ready() requester_tenant_id = (auth_user.tenant_id or "").strip() if not requester_tenant_id: raise HTTPException(status_code=401, detail="Missing tenant_id in JWT claims") if not _session_belongs_to_user(session_id=session_id, user_id=requester_tenant_id): raise HTTPException(status_code=404, detail="Chat session not found") try: # Some versions of Agno support cancel_run success = False if hasattr(gemini_sql_agent, 'cancel_run'): success = gemini_sql_agent.cancel_run(run_id) return {"cancelled": success} except Exception as e: return {"cancelled": False, "error": str(e)} # DEPRECATED FUNCTIONS - Replaced by the /tenant-run API endpoint # The following functions are kept for backward compatibility and local testing only. # For production API usage, use the /tenant-run/{agent_id} endpoint instead. if __name__ == "__main__": import uvicorn host = os.getenv("SQL_AGENT_HOST", "0.0.0.0") port = int(os.getenv("SQL_AGENT_PORT", "5559")) # Override with SQL_AGENT_PORT=8000 for unified print("\n" + "="*80) print("šŸš€ STARTING SQL AGENT OS SERVER (with custom /tenant-run endpoint)") print("="*80) print(f"Host: {host}") print(f"Port: {port}") print(f"Agent ID: {DEFAULT_AGENT_ID}") print(f"AgentOS ID: {DEFAULT_AGENT_OS_ID}") print("="*80 + "\n") print(f"\nšŸŽÆ CUSTOM TENANT ENDPOINT:") print(f" POST http://{host}:{port}/tenant-run/{DEFAULT_AGENT_ID}") print(f"\nšŸ“š STANDARD AGENTOS ENDPOINTS:") print(f" GET http://{host}:{port}/config") print(f" GET http://{host}:{port}/agents") print(f" POST http://{host}:{port}/agents/{DEFAULT_AGENT_ID}/runs") print("="*80 + "\n") # Run with proper streaming settings uvicorn.run( app, host=host, port=port, # Streaming settings - prevent buffering server_header=False, # Disable app level buffering - let streaming work properly interface="auto" )