Update main.py

main.py

@@ -15,32 +15,31 @@ from docx import Document
 import pandas as pd
 import PyPDF2
 
-# Load .env
+# Load environment variables
 load_dotenv()
 
+# Paths for OCR and FFmpeg
+pytesseract.pytesseract.tesseract_cmd = os.getenv("TESSERACT_CMD", "/usr/bin/tesseract")
+os.environ["PATH"] += os.pathsep + os.getenv("FFMPEG_PATH", "/usr/bin")
+
 # Initialize FastAPI
 app = FastAPI()
 
-# Tesseract and FFmpeg setup
-pytesseract.pytesseract.tesseract_cmd = os.getenv("TESSERACT_CMD", "/usr/bin/tesseract")
-ffmpeg_path = os.getenv("FFMPEG_PATH", "/usr/bin")
-os.environ["PATH"] += os.pathsep + ffmpeg_path
+# Load Groq client with key from environment
+client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 
-# Upload directory
 UPLOAD_DIR = "uploaded_files"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
-
-# Groq Client
-client = Groq(api_key=os.getenv("GROQ_API_KEY"))
-
-# File limit
 MAX_FILE_SIZE_MB = 10
 
-
+# --- Helper: Safe Response ---
 def safe_response(text):
-    return text.encode("utf-8", "replace").decode("utf-8")
-
+    try:
+        return str(text).encode("utf-8", "replace").decode("utf-8")
+    except Exception:
+        return "⚠️ Unable to process response."
 
+# --- Extract text from uploaded file ---
 def extract_text_from_file(file_path):
     ext = os.path.splitext(file_path)[-1].lower()
     try:
@@ -48,11 +47,9 @@ def extract_text_from_file(file_path):
             with open(file_path, "r", encoding="utf-8") as f:
                 return f.read()
         elif ext == ".docx":
-            doc = Document(file_path)
-            return "\n".join([p.text for p in doc.paragraphs])
+            return "\n".join([p.text for p in Document(file_path).paragraphs])
         elif ext == ".csv":
-            df = pd.read_csv(file_path)
-            return df.to_string(index=False)
+            return pd.read_csv(file_path).to_string(index=False)
         elif ext == ".pdf":
             with open(file_path, "rb") as f:
                 reader = PyPDF2.PdfReader(f)
@@ -60,40 +57,39 @@ def extract_text_from_file(file_path):
         else:
             return "❌ Unsupported file type."
     except Exception as e:
-        return f"⚠️ Error while reading file: {str(e)}"
-
+        return f"❌ Failed to extract file: {e}"
 
+# --- Endpoint: Chat with File ---
 @app.post("/chat-with-file")
 async def chat_with_file(file: UploadFile = File(...), question: str = Form(...)):
     try:
         contents = await file.read()
         if len(contents) > MAX_FILE_SIZE_MB * 1024 * 1024:
-            return JSONResponse(status_code=400, content={"error": "❌ File too large. Max size is 10MB."})
+            return JSONResponse(content={"answer": "❌ File too large. Max 10MB."})
 
         file_path = os.path.join(UPLOAD_DIR, file.filename)
         with open(file_path, "wb") as f:
             f.write(contents)
 
         file_content = extract_text_from_file(file_path)
-        prompt = f"{file_content}\n\nQuestion: {question}"
 
         response = client.chat.completions.create(
             model="llama3-8b-8192",
             messages=[
-                {"role": "system", "content": "You are a helpful assistant. Use the uploaded file content to answer questions."},
-                {"role": "user", "content": prompt}
+                {"role": "system", "content": "You are a helpful assistant. Use the file content to answer questions."},
+                {"role": "user", "content": f"{file_content}\n\nQuestion: {question}"}
             ]
         )
-        return {"answer": safe_response(response.choices[0].message.content)}
-    except Exception as e:
-        return JSONResponse(status_code=500, content={"error": safe_response(str(e))})
 
+        return JSONResponse(content={"answer": safe_response(response.choices[0].message.content)})
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"answer": f"❌ Error: {e}"})
 
+# --- Endpoint: Chat with URL ---
 class URLQuery(BaseModel):
     url: str
     question: str
 
-
 @app.post("/chat-with-url")
 async def chat_with_url(data: URLQuery):
     try:
@@ -102,31 +98,29 @@ async def chat_with_url(data: URLQuery):
         documents = loader.load()
         web_content = "\n".join([doc.page_content for doc in documents])
 
-        prompt = f"Website Content:\n{web_content}\n\nNow answer this question:\n{data.question}"
-
         response = client.chat.completions.create(
             model="llama3-8b-8192",
             messages=[
-                {"role": "system", "content": "You are a helpful assistant. Use the website content to answer the user's question."},
-                {"role": "user", "content": prompt}
+                {"role": "system", "content": "You are a helpful assistant. Use the website content to answer questions."},
+                {"role": "user", "content": f"Website Content:\n{web_content}\n\nNow answer this:\n{data.question}"}
             ]
         )
-        return {"answer": safe_response(response.choices[0].message.content)}
+        return JSONResponse(content={"answer": safe_response(response.choices[0].message.content)})
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": safe_response(str(e))})
-
+        return JSONResponse(status_code=500, content={"answer": f"❌ Error: {e}"})
 
+# --- Endpoint: OCR from Image ---
 @app.post("/extract-text-from-image")
 async def extract_text_from_image(file: UploadFile = File(...)):
     try:
         contents = await file.read()
         image = Image.open(io.BytesIO(contents))
         text = pytesseract.image_to_string(image)
-        return {"answer": safe_response(text.strip() or "⚠️ No text extracted.")}
+        return JSONResponse(content={"answer": safe_response(text.strip() or "⚠️ No text extracted.")})
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": safe_response(str(e))})
-
+        return JSONResponse(status_code=500, content={"answer": f"❌ Error: {e}"})
 
+# --- Endpoint: Transcribe Audio ---
 @app.post("/transcribe-audio")
 async def transcribe_audio(file: UploadFile = File(...)):
     try:
@@ -137,7 +131,6 @@ async def transcribe_audio(file: UploadFile = File(...)):
 
         model = whisper.load_model("base")
         result = model.transcribe(audio_path)
-        text = result.get("text", "⚠️ No transcript returned.")
-        return {"answer": safe_response(text)}
+        return JSONResponse(content={"answer": safe_response(result.get("text", "⚠️ No transcript returned."))})
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": safe_response(str(e))})
+        return JSONResponse(status_code=500, content={"answer": f"❌ Error: {e}"})