Update main.py

main.py

@@ -15,13 +15,12 @@ from docx import Document
 import pandas as pd
 import PyPDF2
 
-# Load environment variables (GROQ_API_KEY, TESSERACT_CMD, FFMPEG_PATH)
+# Load environment variables
 load_dotenv()
 
-# ✅ Set Tesseract and FFmpeg Paths (Hugging Face-compatible)
+# Set paths for OCR & audio
 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
+os.environ["PATH"] += os.pathsep + os.getenv("FFMPEG_PATH", "/usr/bin")
 
 app = FastAPI()
 client = Groq(api_key=os.getenv("GROQ_API_KEY"))
@@ -31,8 +30,6 @@ os.makedirs(UPLOAD_DIR, exist_ok=True)
 
 MAX_FILE_SIZE_MB = 10
 
-
-# ---------- File Parsing ----------
 def extract_text_from_file(file_path):
     ext = os.path.splitext(file_path)[-1].lower()
     if ext == ".txt":
@@ -40,7 +37,7 @@ def extract_text_from_file(file_path):
             return f.read()
     elif ext == ".docx":
         doc = Document(file_path)
-        return "\n".join([para.text for para in doc.paragraphs])
+        return "\n".join([p.text for p in doc.paragraphs])
     elif ext == ".csv":
         df = pd.read_csv(file_path)
         return df.to_string(index=False)
@@ -48,37 +45,30 @@ def extract_text_from_file(file_path):
         with open(file_path, "rb") as f:
             reader = PyPDF2.PdfReader(f)
             return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
-    else:
-        return "❌ Unsupported file type."
-
+    return "❌ Unsupported file type."
 
-# ---------- File Chat Endpoint ----------
 @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 10MB."})
-
-        file_path = os.path.join(UPLOAD_DIR, file.filename)
-        with open(file_path, "wb") as f:
+        path = os.path.join(UPLOAD_DIR, file.filename)
+        with open(path, "wb") as f:
             f.write(contents)
-
-        file_content = extract_text_from_file(file_path)
+        text = extract_text_from_file(path)
 
         response = client.chat.completions.create(
             model="llama3-8b-8192",
             messages=[
-                {"role": "system", "content": "Use the uploaded file content to answer the user's question."},
-                {"role": "user", "content": f"{file_content}\n\nQuestion: {question}"}
+                {"role": "system", "content": "You are a helpful assistant. Answer using the uploaded file."},
+                {"role": "user", "content": f"{text}\n\nQuestion: {question}"}
             ]
         )
         return {"answer": response.choices[0].message.content}
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
 
-
-# ---------- URL Chat Endpoint ----------
 class URLQuery(BaseModel):
     url: str
     question: str
@@ -87,49 +77,41 @@ class URLQuery(BaseModel):
 async def chat_with_url(data: URLQuery):
     try:
         loader = WebBaseLoader(data.url, header_template={"User-Agent": "Mozilla/5.0"})
-        documents = loader.load()
-        web_content = "\n".join([doc.page_content for doc in documents])
+        docs = loader.load()
+        content = "\n".join([doc.page_content for doc in docs])
 
         response = client.chat.completions.create(
             model="llama3-8b-8192",
             messages=[
-                {"role": "system", "content": "Use website content to answer the user's question."},
-                {"role": "user", "content": f"Website Content:\n{web_content}\n\nQuestion:\n{data.question}"}
+                {"role": "system", "content": "You are a helpful assistant. Answer using the webpage content."},
+                {"role": "user", "content": f"Web Content:\n{content}\n\nQuestion: {data.question}"}
             ]
         )
         return {"answer": response.choices[0].message.content}
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
 
-
-# ---------- Image OCR Endpoint ----------
 @app.post("/extract-text-from-image")
 async def extract_text_from_image(file: UploadFile = File(...)):
     try:
         contents = await file.read()
-
-        # Load and preprocess image
-        image = Image.open(io.BytesIO(contents)).convert("L")  # Grayscale
-        image = image.resize((image.width * 2, image.height * 2))  # Upscale for better OCR
-
-        text = pytesseract.image_to_string(image, lang='eng')
+        image = Image.open(io.BytesIO(contents)).convert("L")
+        image = image.resize((image.width * 2, image.height * 2))
+        text = pytesseract.image_to_string(image, lang="eng")
         return {"answer": text.strip() or "⚠️ No text extracted."}
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
 
-
-
-# ---------- Audio Transcription Endpoint ----------
 @app.post("/transcribe-audio")
 async def transcribe_audio(file: UploadFile = File(...)):
     try:
         contents = await file.read()
-        audio_path = os.path.join(UPLOAD_DIR, file.filename)
-        with open(audio_path, "wb") as f:
+        path = os.path.join(UPLOAD_DIR, file.filename)
+        with open(path, "wb") as f:
             f.write(contents)
 
         model = whisper.load_model("base")
-        result = model.transcribe(audio_path)
-        return {"answer": result["text"] if result.get("text") else "⚠️ No transcript returned."}
+        result = model.transcribe(path)
+        return {"answer": result.get("text", "⚠️ No transcript returned.")}
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})