Update document_generator.py

document_generator.py

@@ -44,6 +44,7 @@ import os
 import json
 import re
 import asyncio
+import time
 from typing import List, Dict, Optional, Any, Callable
 from openai import OpenAI
 import logging
@@ -52,6 +53,9 @@ from fastapi import APIRouter, HTTPException, Request
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from fastapi_cache.decorator import cache
+import psycopg2
+from datetime import datetime
+
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -69,6 +73,28 @@ def log_execution(func: Callable) -> Callable:
             raise
     return wrapper
 
+class DatabaseManager:
+    """Manages database operations."""
+
+    def __init__(self):
+        self.db_params = {
+            "dbname": "postgres",
+            "user": os.environ['SUPABASE_USER'],
+            "password": os.environ['SUPABASE_PASSWORD'],
+            "host": "aws-0-us-west-1.pooler.supabase.com",
+            "port": "5432"
+        }
+
+    @log_execution
+    def update_database(self, user_id: str, user_query: str, response: str) -> None:
+        with psycopg2.connect(**self.db_params) as conn:
+            with conn.cursor() as cur:
+                insert_query = """
+                INSERT INTO ai_document_generator (user_id, user_query, response)
+                VALUES (%s, %s, %s);
+                """
+                cur.execute(insert_query, (user_id, user_query, response))
+
 class AIClient:
     def __init__(self):
         self.client = OpenAI(
@@ -224,7 +250,7 @@ class MarkdownDocumentRequest(BaseModel):
 
 async def generate_document_stream(document_generator: DocumentGenerator, document_outline: Dict, query: str):
     document_generator.document_outline = document_outline
-    
+    db_manager = DatabaseManager()
     overall_objective = query
     document_layout = json.dumps(document_generator.document_outline, indent=2)
 
@@ -276,6 +302,7 @@ async def generate_document_stream(document_generator: DocumentGenerator, docume
         "content": markdown_document,
         "json": document_generator.document_outline
     }, ensure_ascii=False) + "\n"
+    db_manager.update_database("elevatics", query, markdown_document)
 
 @cache(expire=600*24*7)
 @router.post("/generate-document/json", response_model=JsonDocumentResponse)