feat(chatDraft): rewrite with AI in chat draft

Author: ArslanSaleem

backend/app/api/v1/chat.py

@@ -10,11 +10,11 @@
     project_repository,
     user_repository,
 )
-from app.requests import chat_query
+from app.requests import chat_query, request_draft_with_ai
 from app.utils import clean_text, find_following_sentence_ending, find_sentence_endings
 from app.vectorstore.chroma import ChromaDB
 from fastapi import APIRouter, Depends, HTTPException
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
 
 chat_router = APIRouter()
@@ -24,6 +24,10 @@ class ChatRequest(BaseModel):
     conversation_id: Optional[str] = None
     query: str
 
+class DraftRequest(BaseModel):
+    content: str = Field(..., min_length=1, description="Content cannot be empty")
+    prompt: str = Field(..., min_length=1, description="Prompt cannot be empty")
+
 
 logger = Logger()
 
@@ -237,3 +241,35 @@ def chat_status(project_id: int, db: Session = Depends(get_db)):
             status_code=400,
             detail="Unable to process the chat query. Please try again.",
         )
+
+@chat_router.post("/draft", status_code=200)
+def draft_with_ai(draft_request: DraftRequest, db: Session = Depends(get_db)):
+    try:
+
+        users = user_repository.get_users(db, n=1)
+
+        if not users:
+            raise HTTPException(status_code=404, detail="No User Exists!")
+
+        api_key = user_repository.get_user_api_key(db, users[0].id)
+
+        if not api_key:
+            raise HTTPException(status_code=404, detail="API Key not found!")
+
+        response = request_draft_with_ai(api_key.key, draft_request.model_dump_json())
+
+        return {
+            "status": "success",
+            "message": "Chat message successfully generated.",
+            "data": {"response": response["response"]},
+        }
+
+    except HTTPException:
+        raise
+
+    except Exception:
+        logger.error(traceback.format_exc())
+        raise HTTPException(
+            status_code=400,
+            detail="Unable to process the chat query. Please try again.",
+        )

backend/app/requests/__init__.py

@@ -220,3 +220,27 @@ def get_user_usage_data(api_token: str):
     except requests.exceptions.JSONDecodeError:
         logger.error(f"Invalid JSON response from API server: {response.text}")
         raise Exception("Invalid JSON response")
+
+
+def request_draft_with_ai(api_token: str, draft_request: dict) -> dict:
+    # Prepare the headers with the Bearer token
+    headers = {"x-authorization": f"Bearer {api_token}"}
+    # Send the request
+    response = requests.post(
+        f"{settings.pandaetl_server_url}/v1/draft",
+        data=draft_request,
+        headers=headers,
+        timeout=360,
+    )
+
+    try:
+        if response.status_code not in [200, 201]:
+            logger.error(
+                f"Failed to draft with AI. It returned {response.status_code} code: {response.text}"
+            )
+            raise Exception(response.text)
+
+        return response.json()
+    except requests.exceptions.JSONDecodeError:
+        logger.error(f"Invalid JSON response from API server: {response.text}")
+        raise Exception("Invalid JSON response")

frontend/src/components/ChatDraftDrawer.tsx

@@ -1,9 +1,12 @@
 "use client";
-import React, { useEffect, useRef } from "react";
+import React, { useEffect, useRef, useState } from "react";
 import Drawer from "./ui/Drawer";
 import { Button } from "./ui/Button";
 import ReactQuill from "react-quill";
-import { BookTextIcon } from "lucide-react";
+import { BookTextIcon, Check, Loader2, X } from "lucide-react";
+import { Textarea } from "./ui/Textarea";
+import { draft_with_ai } from "@/services/chat";
+import toast from "react-hot-toast";
 
 interface IProps {
   draft: string;
@@ -48,6 +51,10 @@ const ChatDraftDrawer = ({
   onCancel,
 }: IProps) => {
   const quillRef = useRef<ReactQuill | null>(null);
+  const [step, setStep] = useState<number>(0);
+  const [userInput, setUserInput] = useState<string>("");
+  const [aiDraft, setAIDraft] = useState<string>("");
+  const [loadingAIDraft, setLoadingAIDraft] = useState<boolean>(false);
 
   useEffect(() => {
     if (quillRef.current) {
@@ -59,28 +66,131 @@ const ChatDraftDrawer = ({
     }
   }, [draft]);
 
+  const handleUserInputChange = (
+    event: React.ChangeEvent<HTMLTextAreaElement>
+  ) => {
+    setUserInput(event.target.value);
+  };
+
+  const handleUserInputKeyPress = async (
+    event: React.KeyboardEvent<HTMLTextAreaElement>
+  ) => {
+    if (event.key === "Enter" && userInput.trim() !== "") {
+      event.preventDefault();
+      try {
+        if (userInput.length === 0) {
+          toast.error("Please provide the prompt and try again!");
+          return;
+        }
+        setLoadingAIDraft(true);
+        const data = await draft_with_ai({ content: draft, prompt: userInput });
+        setAIDraft(data.response);
+        setUserInput("");
+        setStep(2);
+        setLoadingAIDraft(false);
+      } catch (error) {
+        console.error(error);
+        toast.error("Failed to draft with AI. Try again!");
+        setLoadingAIDraft(false);
+      }
+    }
+  };
+
   return (
     <Drawer isOpen={isOpen} onClose={onCancel} title="Draft Chat">
       <div className="flex flex-col h-full">
-        <ReactQuill
-          ref={quillRef}
-          theme="snow"
-          value={draft}
-          onChange={onSubmit}
-          modules={quill_modules}
-          formats={quill_formats}
-        />
-        <div className="sticky bottom-0 bg-white pb-4">
-          <div className="flex gap-2">
-            <Button
-              // onClick={onSubmit}
-              className="mt-4 px-4 py-2 bg-primary text-white rounded hover:bg-primary-dark"
-            >
-              <BookTextIcon className="inline-block mr-2" size={16} />
-              Rewrite with AI
-            </Button>
+        {(step === 0 || step === 1) && (
+          <>
+            <ReactQuill
+              ref={quillRef}
+              theme="snow"
+              value={draft}
+              onChange={onSubmit}
+              modules={quill_modules}
+              formats={quill_formats}
+            />
+
+            <div className="sticky bottom-0 bg-white pb-4 pt-4">
+              <Button
+                onClick={() => {
+                  setStep(1);
+                }}
+                disabled={draft.length == 0}
+                className="px-4 bg-primary text-white rounded hover:bg-primary-dark"
+              >
+                <BookTextIcon className="inline-block mr-2" size={16} />
+                Rewrite with AI
+              </Button>
+            </div>
+          </>
+        )}
+
+        {step === 2 && (
+          <>
+            <ReactQuill
+              ref={quillRef}
+              theme="snow"
+              value={aiDraft}
+              readOnly={true}
+              modules={{ toolbar: false }}
+            />
+
+            <div className="sticky bottom-0 bg-white pb-4">
+              <div className="flex gap-2">
+                <Button
+                  onClick={() => setStep(0)}
+                  className="mt-4 px-4 py-2 bg-red-600 text-white rounded hover:bg-red-800"
+                >
+                  <X className="inline-block mr-2" size={16} />
+                  Cancel
+                </Button>
+                <Button
+                  onClick={() => {
+                    onSubmit(aiDraft);
+                    setStep(0);
+                  }}
+                  className="mt-4 px-4 py-2 bg-green-600 text-white rounded hover:bg-green-800"
+                >
+                  <Check className="inline-block mr-2" size={16} />
+                  Accept
+                </Button>
+                <Button
+                  onClick={() => setStep(1)}
+                  className="mt-4 px-4 py-2 bg-primary text-white rounded hover:bg-primary-dark"
+                >
+                  <BookTextIcon className="inline-block mr-2" size={16} />
+                  Rewrite
+                </Button>
+              </div>
+            </div>
+          </>
+        )}
+        {/* Centered overlay input for step 1 */}
+        {step === 1 && (
+          <div className="absolute inset-0 flex items-center justify-center z-50 bg-opacity-75 bg-gray-800">
+            <div className="bg-white p-6 rounded-lg shadow-lg max-w-lg w-full text-center">
+              {loadingAIDraft ? (
+                <Loader2 className="mx-auto my-4 h-8 w-8 animate-spin text-gray-500" />
+              ) : (
+                <>
+                  <Textarea
+                    className="w-full p-2 border border-gray-300 rounded mb-4"
+                    placeholder="Write prompt to edit content and press enter..."
+                    value={userInput}
+                    onChange={handleUserInputChange}
+                    onKeyDown={handleUserInputKeyPress}
+                  />
+                  <Button
+                    onClick={() => setStep(0)}
+                    className="text-sm text-gray-500 hover:text-gray-700"
+                  >
+                    Close
+                  </Button>
+                </>
+              )}
+            </div>
           </div>
-        </div>
+        )}
       </div>
     </Drawer>
   );

frontend/src/interfaces/chat.ts

@@ -27,3 +27,8 @@ export interface ChatReferences {
   start: number;
   end: number;
 }
+
+export interface ChatDraftRequest {
+  content: string;
+  prompt: string;
+}

frontend/src/services/chat.tsx

@@ -1,6 +1,7 @@
 import axios from "axios";
 import { GetRequest, PostRequest } from "@/lib/requests";
 import {
+  ChatDraftRequest,
   ChatRequest,
   ChatResponse,
   ChatStatusResponse,
@@ -51,3 +52,17 @@ export const chatStatus = async (projectId: string) => {
     }
   }
 };
+
+export const draft_with_ai = async (data: ChatDraftRequest) => {
+  try {
+    const response = await PostRequest<ChatResponse>(
+      `${chatApiUrl}/draft`,
+      { ...data },
+      {},
+      300000
+    );
+    return response.data.data;
+  } catch (error) {
+    throw error;
+  }
+};