Merge pull request #19 from tinybirdco/explorer-02

Explorer 02

Author: alrocar

dashboard/ai-analytics/src/app/api/extract-cost-parameters/route.ts

@@ -25,6 +25,73 @@ const costParametersSchema = z.object({
 // Create a type from the schema
 type CostParameters = z.infer<typeof costParametersSchema>;
 
+// Add a new function to fetch available dimensions
+const fetchAvailableDimensions = async () => {
+  const TINYBIRD_API_URL = process.env.NEXT_PUBLIC_TINYBIRD_API_URL || 'http://localhost:7181';
+  const TINYBIRD_API_KEY = process.env.NEXT_PUBLIC_TINYBIRD_API_KEY;
+  
+  if (!TINYBIRD_API_KEY) {
+    console.error('No Tinybird API key available');
+    return null;
+  }
+  
+  try {
+    // SQL query to get all unique values for each dimension
+    const query = `
+      SELECT
+        groupUniqArray(organization) as organizations,
+        groupUniqArray(project) as projects,
+        groupUniqArray(environment) as environments,
+        groupUniqArray(model) as models,
+        groupUniqArray(provider) as providers
+      FROM llm_events FORMAT JSON
+    `;
+    
+    // URL encode the query
+    const encodedQuery = encodeURIComponent(query);
+    const url = `${TINYBIRD_API_URL}/v0/sql?q=${encodedQuery}`;
+    
+    console.log('Fetching available dimensions from:', url);
+    
+    const response = await fetch(url, {
+      headers: {
+        Authorization: `Bearer ${TINYBIRD_API_KEY}`,
+      },
+    });
+    
+    if (!response.ok) {
+      const error = await response.text();
+      console.error('Error fetching dimensions:', error);
+      throw new Error('Network response was not ok');
+    }
+    
+    const data = await response.json();
+    console.log('Available dimensions:', data);
+    return data;
+  } catch (error) {
+    console.error('Error fetching dimensions:', error);
+    return null;
+  }
+};
+
+// Format today's date with time in yyyy-MM-dd HH:mm:ss format
+const formatDate = (date: Date): string => {
+  const pad = (num: number): string => num.toString().padStart(2, '0');
+  
+  const year = date.getFullYear();
+  const month = pad(date.getMonth() + 1); // getMonth() is 0-indexed
+  const day = pad(date.getDate());
+  const hours = pad(date.getHours());
+  const minutes = pad(date.getMinutes());
+  const seconds = pad(date.getSeconds());
+  
+  return `${year}-${month}-${day} ${hours}:${minutes}:${seconds}`;
+};
+
+const today = formatDate(new Date());
+const todayDateOnly = today.split(' ')[0]; // Just the date part: yyyy-MM-dd
+
+// Update the POST function to properly map meta with data
 export async function POST(req: Request) {
   try {
     const { query } = await req.json();
@@ -33,13 +100,39 @@ export async function POST(req: Request) {
       return NextResponse.json({ error: 'Query is required' }, { status: 400 });
     }
 
-    // Get today's date for the prompt
-    const today = new Date().toISOString().split('T')[0]; // Format: YYYY-MM-DD
+    // Fetch pipe definition and available dimensions in parallel
+    const [pipeDefinition, availableDimensions] = await Promise.all([
+      fetchPipeDefinition(),
+      fetchAvailableDimensions()
+    ]);
+    
+    console.log('Pipe definition:', pipeDefinition);
+    console.log('Available dimensions:', availableDimensions);
 
+    // Extract dimension values for the system prompt
+    const dimensionValues: Record<string, { type: string, values: string[] }> = {};
+    
+    // Map meta with data for a more structured representation
+    if (availableDimensions && availableDimensions.meta && availableDimensions.data && availableDimensions.data.length > 0) {
+      const metaInfo = availableDimensions.meta;
+      const dataValues = availableDimensions.data[0];
+      
+      // Create a structured object with meta information and data values
+      metaInfo.forEach((meta: { name: string, type: string }) => {
+        dimensionValues[meta.name] = {
+          type: meta.type,
+          values: dataValues[meta.name] || []
+        };
+      });
+    }
+    
+    console.log('Structured dimension values:', dimensionValues);
+    
+    // Update system prompt to include formatted timestamp
     const systemPromptText = `
-      You are a cost prediction parameter extractor. Extract parameters from natural language queries about AI model cost predictions.
+      You are an LLM cost calculator parameter extractor. Extract parameters from natural language queries about AI model cost predictions.
       
-      Today's date is ${today}.
+      Today's date and time is ${today}.
       
       Return values for these parameters:
       - model: The AI model mentioned (e.g., "gpt-4", "claude-sonnet")
@@ -49,15 +142,23 @@ export async function POST(req: Request) {
       - timeframe: Time period for analysis (e.g., "last month", "last week", "last 3 months", "last year")
       - volumeChange: Any volume change percentage mentioned (can be positive or negative)
       - start_date: The start date in YYYY-MM-DD format based on the timeframe
-      - end_date: The end date in YYYY-MM-DD format (usually today: ${today})
+      - end_date: The end date in YYYY-MM-DD format (usually today: ${todayDateOnly})
       - group_by: If the user wants to group by a specific dimension (e.g., "model", "provider", "organization", "project", "environment", "user")
       
-      Filter parameters (extract if mentioned):
-      - organization: Organization name to filter by (e.g., "quantum_systems", "acme_corp")
-      - project: Project name to filter by (e.g., "chatbot", "recommendation_engine")
-      - environment: Environment to filter by (e.g., "production", "staging", "development")
-      - provider: Provider to filter by (e.g., "OpenAI", "Anthropic", "Cohere")
-      - user: User to filter by (e.g., "john.doe", "api_user")
+      Also extract filter parameters if mentioned:
+      ${availableDimensions?.meta?.map((meta: { name: string, type: string }) => {
+        const name = meta.name;
+        // Convert to singular form for the parameter name (remove trailing 's')
+        const paramName = name.endsWith('s') ? name.slice(0, -1) : name;
+        return `- ${paramName}: ${name.charAt(0).toUpperCase() + name.slice(1)} to filter by`;
+      }).join('\n  ') || ''}
+      
+      These are the available values in the database, extract them if mentioned no matter if the parameter name is provided or not:
+      ${availableDimensions?.meta?.map((meta: { name: string, type: string }) => {
+        const name = meta.name;
+        const values = dimensionValues[name]?.values || [];
+        return `- ${name.charAt(0).toUpperCase() + name.slice(1)}: ${JSON.stringify(values)}`;
+      }).join('\n  ') || ''}
       
       For timeframes, calculate the appropriate start_date:
       - "last week" = 7 days before today
@@ -69,12 +170,7 @@ export async function POST(req: Request) {
       Always include start_date and end_date based on the timeframe (default to "last month" if not specified).
       
       Context-based filter extraction rules:
-      1. When a model name is mentioned (e.g., "gpt-4", "claude-3-opus"), set model filter to that value
-      2. When a provider name is mentioned (e.g., "OpenAI", "Anthropic"), set provider filter to that value.
-      3. When an environment is mentioned (e.g., "production", "staging", "development"), set environment filter
-      4. When a project or organization name is mentioned, set the appropriate filter
-      5. When a username is mentioned, set the user filter
-      6. Never use a model name as a provider name and vice-versa, model=anthropic is invalid, provider=anthropic is valid.
+        - Only use values that exist in the database (see "Available values" above).
       
       Look for phrases like "filter by", "for", "in", "with", etc. to identify filter parameters.
       Examples:
@@ -127,7 +223,9 @@ export async function POST(req: Request) {
       project: extractedParams.project || null,
       environment: extractedParams.environment || null,
       provider: extractedParams.provider || null,
-      user: extractedParams.user || null
+      user: extractedParams.user || null,
+      pipeDefinition: pipeDefinition,
+      availableDimensions: dimensionValues
     };
 
     return NextResponse.json(processedResult);
@@ -163,5 +261,41 @@ function getDefaultStartDate(timeframe: string): string {
     startDate.setMonth(now.getMonth() - 1);
   }
 
-  return startDate.toISOString().split('T')[0];
-}
+  // Format the date with the current time
+  return formatDate(startDate);
+}
+
+// Fetch the llm_usage pipe definition
+const fetchPipeDefinition = async () => {
+  const TINYBIRD_API_URL = process.env.NEXT_PUBLIC_TINYBIRD_API_URL || 'http://localhost:7181';
+  const TINYBIRD_API_KEY = process.env.NEXT_PUBLIC_TINYBIRD_API_KEY;
+  
+  if (!TINYBIRD_API_KEY) {
+    console.error('No Tinybird API key available');
+    return null;
+  }
+  
+  try {
+    const url = `${TINYBIRD_API_URL}/v0/pipes/llm_usage`;
+    console.log('Fetching pipe definition from:', url);
+    
+    const response = await fetch(url, {
+      headers: {
+        Authorization: `Bearer ${TINYBIRD_API_KEY}`,
+      },
+    });
+    
+    if (!response.ok) {
+      const error = await response.text();
+      console.error('Error fetching pipe definition:', error);
+      throw new Error('Network response was not ok');
+    }
+    
+    const data = await response.json();
+    console.log('Pipe definition:', data);
+    return data;
+  } catch (error) {
+    console.error('Error fetching pipe definition:', error);
+    return null;
+  }
+};