Update observability.py

observability.py

@@ -110,67 +110,160 @@ class LLMObservabilityManager:
             column_names = [description[0] for description in cursor.description]
             return [dict(zip(column_names, row)) for row in rows]
 
-    ## OBSERVABILITY
-from uuid import uuid4
-import csv
-from io import StringIO
-from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel
-from starlette.responses import StreamingResponse
-
-
-
-router = APIRouter( 
-    prefix="/observability",
-    tags=["observability"]
-)
-
-class ObservationResponse(BaseModel):
-    observations: List[Dict]
-    
-def create_csv_response(observations: List[Dict]) -> StreamingResponse:
-    def iter_csv(data):
-        output = StringIO()
-        writer = csv.DictWriter(output, fieldnames=data[0].keys() if data else [])
-        writer.writeheader()
-        for row in data:
-            writer.writerow(row)
-        output.seek(0)
-        yield output.read()
-
-    headers = {
-        'Content-Disposition': 'attachment; filename="observations.csv"'
-    }
-    return StreamingResponse(iter_csv(observations), media_type="text/csv", headers=headers)
-    
-
-@router.get("/last-observations/{limit}")
-async def get_last_observations(limit: int = 10, format: str = "json"):
-    observability_manager = LLMObservabilityManager()
-    
-    try:
-        # Get all observations, sorted by created_at in descending order
-        all_observations = observability_manager.get_observations()
-        all_observations.sort(key=lambda x: x['created_at'], reverse=True)
+    def get_dashboard_statistics(self, days: Optional[int] = None, time_series_interval: str = 'day') -> Dict[str, Any]:
+        """
+        Get statistical metrics for LLM usage dashboard with time series data.
         
-        # Get the last conversation_id
-        if all_observations:
-            last_conversation_id = all_observations[0]['conversation_id']
+        Args:
+            days (int, optional): Number of days to look back. If None, returns all-time statistics
+            time_series_interval (str): Interval for time series data ('hour', 'day', 'week', 'month')
             
-            # Filter observations for the last conversation
-            last_conversation_observations = [
-                obs for obs in all_observations
-                if obs['conversation_id'] == last_conversation_id
-            ][:limit]
+        Returns:
+            Dict containing dashboard statistics and time series data
+        """
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
             
-            if format.lower() == "csv":
-                return create_csv_response(last_conversation_observations)
-            else:
-                return ObservationResponse(observations=last_conversation_observations)
-        else:
-            if format.lower() == "csv":
-                return create_csv_response([])
+            # Build time filter
+            time_filter = ""
+            if days is not None:
+                time_filter = f"WHERE created_at >= datetime('now', '-{days} days')"
+            
+            # Get general statistics
+            cursor.execute(f"""
+                SELECT 
+                    COUNT(*) as total_requests,
+                    COUNT(DISTINCT conversation_id) as unique_conversations,
+                    COUNT(DISTINCT user) as unique_users,
+                    SUM(total_tokens) as total_tokens,
+                    SUM(cost) as total_cost,
+                    AVG(latency) as avg_latency,
+                    SUM(CASE WHEN status = 'error' THEN 1 ELSE 0 END) as error_count
+                FROM llm_observations
+                {time_filter}
+            """)
+            general_stats = dict(zip([col[0] for col in cursor.description], cursor.fetchone()))
+            
+            # Get model distribution
+            cursor.execute(f"""
+                SELECT model, COUNT(*) as count
+                FROM llm_observations
+                {time_filter}
+                GROUP BY model
+                ORDER BY count DESC
+            """)
+            model_distribution = {row[0]: row[1] for row in cursor.fetchall()}
+            
+            # Get average tokens per request
+            cursor.execute(f"""
+                SELECT 
+                    AVG(prompt_tokens) as avg_prompt_tokens,
+                    AVG(completion_tokens) as avg_completion_tokens
+                FROM llm_observations
+                {time_filter}
+            """)
+            token_averages = dict(zip([col[0] for col in cursor.description], cursor.fetchone()))
+            
+            # Get top users by request count
+            cursor.execute(f"""
+                SELECT user, COUNT(*) as request_count, 
+                       SUM(total_tokens) as total_tokens,
+                       SUM(cost) as total_cost
+                FROM llm_observations
+                {time_filter}
+                GROUP BY user
+                ORDER BY request_count DESC
+                LIMIT 5
+            """)
+            top_users = [
+                {
+                    "user": row[0],
+                    "request_count": row[1],
+                    "total_tokens": row[2],
+                    "total_cost": round(row[3], 2)
+                }
+                for row in cursor.fetchall()
+            ]
+    
+            # Get time series data
+            time_series_format = {
+                'hour': "%Y-%m-%d %H:00:00",
+                'day': "%Y-%m-%d",
+                'week': "%Y-%W",
+                'month': "%Y-%m"
+            }
+            
+            format_string = time_series_format[time_series_interval]
+            
+            cursor.execute(f"""
+                SELECT 
+                    strftime('{format_string}', created_at) as time_bucket,
+                    COUNT(*) as request_count,
+                    SUM(total_tokens) as total_tokens,
+                    SUM(cost) as total_cost,
+                    AVG(latency) as avg_latency,
+                    COUNT(DISTINCT user) as unique_users,
+                    SUM(CASE WHEN status = 'error' THEN 1 ELSE 0 END) as error_count
+                FROM llm_observations
+                {time_filter}
+                GROUP BY time_bucket
+                ORDER BY time_bucket
+            """)
+            
+            time_series = [
+                {
+                    "timestamp": row[0],
+                    "request_count": row[1],
+                    "total_tokens": row[2],
+                    "total_cost": round(row[3], 2),
+                    "avg_latency": round(row[4], 2),
+                    "unique_users": row[5],
+                    "error_count": row[6]
+                }
+                for row in cursor.fetchall()
+            ]
+    
+            # Calculate usage trends (percentage change)
+            if len(time_series) >= 2:
+                current = time_series[-1]
+                previous = time_series[-2]
+                trends = {
+                    "request_trend": calculate_percentage_change(
+                        previous["request_count"], current["request_count"]),
+                    "cost_trend": calculate_percentage_change(
+                        previous["total_cost"], current["total_cost"]),
+                    "token_trend": calculate_percentage_change(
+                        previous["total_tokens"], current["total_tokens"])
+                }
             else:
-                return ObservationResponse(observations=[])
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to retrieve observations: {str(e)}")
+                trends = {
+                    "request_trend": 0,
+                    "cost_trend": 0,
+                    "token_trend": 0
+                }
+    
+            return {
+                "general_stats": {
+                    "total_requests": general_stats["total_requests"],
+                    "unique_conversations": general_stats["unique_conversations"],
+                    "unique_users": general_stats["unique_users"],
+                    "total_tokens": general_stats["total_tokens"],
+                    "total_cost": round(general_stats["total_cost"], 2),
+                    "avg_latency": round(general_stats["avg_latency"], 2),
+                    "error_rate": round(general_stats["error_count"] / general_stats["total_requests"] * 100, 2)
+                },
+                "model_distribution": model_distribution,
+                "token_metrics": {
+                    "avg_prompt_tokens": round(token_averages["avg_prompt_tokens"], 2),
+                    "avg_completion_tokens": round(token_averages["avg_completion_tokens"], 2)
+                },
+                "top_users": top_users,
+                "time_series": time_series,
+                "trends": trends
+            }
+    
+    def calculate_percentage_change(old_value: float, new_value: float) -> float:
+        """Calculate percentage change between two values."""
+        if old_value == 0:
+            return 100 if new_value > 0 else 0
+        return round(((new_value - old_value) / old_value) * 100, 2)