Update app.py

app.py

@@ -5,6 +5,8 @@ import numpy as np
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
 from typing import List, Tuple
+from fastapi import FastAPI, Query
+import uvicorn
 
 # Default settings
 class ChatConfig:
@@ -40,39 +42,29 @@ def search_relevant_text(query):
     _, closest_idx = index.search(np.array(query_embedding, dtype=np.float32), k=3)
     return "\n".join([documents[i] for i in closest_idx[0]])
 
-def generate_response(
-    message: str,
-    history: List[Tuple[str, str]],
-    system_message: str = ChatConfig.DEFAULT_SYSTEM_MSG,
-    max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
-    temperature: float = ChatConfig.DEFAULT_TEMP,
-    top_p: float = ChatConfig.DEFAULT_TOP_P
-) -> str:
+def generate_response_sync(message: str) -> str:
+    """Generates response synchronously for FastAPI"""
     if not documents:
         return "Please upload a PDF first."
 
     context = search_relevant_text(message)  # Get relevant content from PDF
-
-    messages = [{"role": "system", "content": system_message}]
-    for user_msg, bot_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
-
-    messages.append({"role": "user", "content": f"Context: {context}\nQuestion: {message}"})
+    messages = [
+        {"role": "system", "content": ChatConfig.DEFAULT_SYSTEM_MSG},
+        {"role": "user", "content": f"Context: {context}\nQuestion: {message}"}
+    ]
 
     response = ""
     for chunk in client.chat_completion(
         messages,
-        max_tokens=max_tokens,
+        max_tokens=ChatConfig.DEFAULT_MAX_TOKENS,
         stream=True,
-        temperature=temperature,
-        top_p=top_p,
+        temperature=ChatConfig.DEFAULT_TEMP,
+        top_p=ChatConfig.DEFAULT_TOP_P,
     ):
         token = chunk.choices[0].delta.content or ""
         response += token
-        yield response
+
+    return response
 
 def handle_upload(pdf_file):
     """Handles PDF upload and creates vector DB"""
@@ -100,13 +92,31 @@ def create_interface() -> gr.Blocks:
 
         # Chat function
         send_button.click(
-            generate_response, 
-            inputs=[user_input, chatbot], 
+            generate_response_sync, 
+            inputs=[user_input], 
             outputs=[output]
         )
 
     return interface
 
+# FastAPI Integration
+app = FastAPI()
+
+@app.get("/chat")
+def chat_with_pdf(msg: str = Query(..., title="User Message")):
+    """API endpoint to receive a message and return AI response"""
+    response = generate_response_sync(msg)
+    return {"response": response}
+
 if __name__ == "__main__":
-    app = create_interface()
-    app.launch()
+    import threading
+
+    # Start Gradio UI in a separate thread
+    def run_gradio():
+        gradio_app = create_interface()
+        gradio_app.launch(server_name="0.0.0.0", server_port=7860, share=True)
+
+    threading.Thread(target=run_gradio).start()
+
+    # Start FastAPI
+    uvicorn.run(app, host="0.0.0.0", port=8000)