Update qwen_app.py

qwen_app.py

@@ -1,153 +1,212 @@
-import os
-from typing import List, Any
-from chainlit.types import AskFileResponse
-import tempfile
-import shutil
-
-# Text processing
-from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.document_loaders import TextLoader, PyPDFLoader
-from langchain.docstore.document import Document
-
-# Prompt templates
-from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
-
-# Embeddings + VectorDB
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-
-# OpenRouter (Qwen via OpenAI-compatible API)
-from langchain_openai import ChatOpenAI
-
-# UI framework
-import chainlit as cl
-
-
-# -------------------------
-# API Key Setup
-# -------------------------
-# Make sure you export in Colab / Terminal before running:
-#   os.environ["OPENROUTER_API_KEY"] = "your_api_key_here"
-#   os.environ["OPENAI_API_BASE"] = "https://openrouter.ai/api/v1"
-
-
-# -------------------------
-# File processing
-# -------------------------
-text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-
-def process_file(file: AskFileResponse) -> List[Document]:
-    """Load and split PDF or TXT into LangChain Documents."""
-    suffix = f".{file.name.split('.')[-1]}"
-    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
-        shutil.copyfile(file.path, temp_file.name)
-
-        if file.name.lower().endswith(".pdf"):
-            loader = PyPDFLoader(temp_file.name)
-        else:
-            loader = TextLoader(temp_file.name)
-
-        try:
-            docs = loader.load()
-            texts = text_splitter.split_documents(docs)
-            return texts
-        finally:
-            try:
-                os.unlink(temp_file.name)
-            except Exception as e:
-                print(f"Cleanup error: {e}")
-
-
-# -------------------------
-# Retrieval QA Pipeline
-# -------------------------
-class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: Any, vectorstore: FAISS) -> None:
-        self.llm = llm
-        self.vectorstore = vectorstore
-
-        # Prompt definition
-        system_template = (
-            "You are a helpful assistant. "
-            "Use the following context to answer a user's question. "
-            "If the context does not contain the answer, reply with 'I don't know'."
-        )
-        self.prompt = ChatPromptTemplate.from_messages([
-            SystemMessagePromptTemplate.from_template(system_template),
-            HumanMessagePromptTemplate.from_template("Context:\n{context}\n\nQuestion:\n{question}")
-        ])
-
-    async def arun_pipeline(self, user_query: str):
-        # Retrieve documents
-        docs = self.vectorstore.similarity_search(user_query, k=4)
-        context_text = "\n".join([doc.page_content for doc in docs])
-
-        # Format the prompt
-        messages = self.prompt.format_messages(context=context_text, question=user_query)
-
-        # Stream response from Qwen
-        async def generate_response():
-            async for chunk in self.llm.astream(messages):
-                yield chunk.content if chunk.content else ""
-
-        return {"response": generate_response(), "context": docs}
-
-
-# -------------------------
-# Chainlit Handlers
-# -------------------------
-@cl.on_chat_start
-async def on_chat_start():
-    files = None
-
-    # Wait for user file
-    while files is None:
-        files = await cl.AskFileMessage(
-            content="Please upload a Text or PDF file to begin!",
-            accept=["text/plain", "application/pdf"],
-            max_size_mb=5,
-            timeout=180,
-        ).send()
-
-    file = files[0]
-    msg = cl.Message(content=f"Processing `{file.name}`...")
-    await msg.send()
-
-    # Load & process file
-    texts = process_file(file)
-    print(f"Processing {len(texts)} chunks")
-
-    # Create embeddings + vectorstore
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    vectorstore = FAISS.from_documents(texts, embeddings)
-
-    # Initialize Qwen via OpenRouter
-    chat_llm = ChatOpenAI(
-        model="qwen/qwen2.5-vl-72b-instruct",  # ✅ you can swap with qwen-3 when available
-        streaming=True,
-        temperature=0,
-        max_tokens=1024,
-        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://openrouter.ai/api/v1"),
-        openai_api_key= "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"
-    )
-
-    # Create retrieval pipeline
-    retrieval_qa = RetrievalAugmentedQAPipeline(llm=chat_llm, vectorstore=vectorstore)
-
-    msg.content = f"Processing `{file.name}` done ✅. You can now ask questions!"
-    await msg.update()
-
-    cl.user_session.set("chain", retrieval_qa)
-
-
-@cl.on_message
-async def main(message: cl.Message):
-    chain = cl.user_session.get("chain")
-
-    msg = cl.Message(content="")
-    result = await chain.arun_pipeline(message.content)
-
-    async for stream_resp in result["response"]:
-        await msg.stream_token(stream_resp)
-
-    await msg.send()
+import os
+import tempfile
+import shutil
+from typing import List, Any
+
+from chainlit.types import AskFileResponse
+import chainlit as cl
+
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.document_loaders import TextLoader, PyPDFLoader
+from langchain.docstore.document import Document
+
+from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_openai import ChatOpenAI
+
+from fastapi import FastAPI, UploadFile, Form
+from fastapi.responses import JSONResponse
+import uvicorn
+import requests
+
+
+# -------------------------
+# File processing
+# -------------------------
+text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+def process_file_path(file_path: str) -> List[Document]:
+    """Load and split PDF or TXT into LangChain Documents."""
+    if file_path.lower().endswith(".pdf"):
+        loader = PyPDFLoader(file_path)
+    else:
+        loader = TextLoader(file_path)
+
+    docs = loader.load()
+    return text_splitter.split_documents(docs)
+
+
+# -------------------------
+# Retrieval QA Pipeline
+# -------------------------
+class RetrievalAugmentedQAPipeline:
+    def __init__(self, llm: Any, vectorstore: FAISS) -> None:
+        self.llm = llm
+        self.vectorstore = vectorstore
+
+        system_template = (
+            "You are a helpful assistant. "
+            "Use the following context to answer a user's question. "
+            "If the context does not contain the answer, reply with 'I don't know'."
+        )
+
+        self.prompt = ChatPromptTemplate.from_messages([
+            SystemMessagePromptTemplate.from_template(system_template),
+            HumanMessagePromptTemplate.from_template("Context:\n{context}\n\nQuestion:\n{question}")
+        ])
+
+    async def arun_pipeline(self, user_query: str):
+        docs = self.vectorstore.similarity_search(user_query, k=4)
+        context_text = "\n".join([doc.page_content for doc in docs])
+        messages = self.prompt.format_messages(context=context_text, question=user_query)
+
+        async def generate_response():
+            async for chunk in self.llm.astream(messages):
+                yield chunk.content if chunk.content else ""
+
+        return {"response": generate_response(), "context": docs}
+
+
+# -------------------------
+# Chainlit Handlers (UI)
+# -------------------------
+@cl.on_chat_start
+async def on_chat_start():
+    files = None
+    while files is None:
+        files = await cl.AskFileMessage(
+            content="Please upload a Text or PDF file to begin!",
+            accept=["text/plain", "application/pdf"],
+            max_size_mb=5,
+            timeout=180,
+        ).send()
+
+    file = files[0]
+    msg = cl.Message(content=f"Processing `{file.name}`...")
+    await msg.send()
+
+    texts = process_file_path(file.path)
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_documents(texts, embeddings)
+
+    chat_llm = ChatOpenAI(
+        model="qwen/qwen2.5-vl-72b-instruct",
+        streaming=True,
+        temperature=0,
+        max_tokens=1024,
+        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://openrouter.ai/api/v1"),
+        openai_api_key= "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"
+    )
+
+    retrieval_qa = RetrievalAugmentedQAPipeline(llm=chat_llm, vectorstore=vectorstore)
+    msg.content = f"Processing `{file.name}` done ✅. You can now ask questions!"
+    await msg.update()
+    cl.user_session.set("chain", retrieval_qa)
+
+
+@cl.on_message
+async def main(message: cl.Message):
+    chain = cl.user_session.get("chain")
+    msg = cl.Message(content="")
+    result = await chain.arun_pipeline(message.content)
+
+    async for stream_resp in result["response"]:
+        await msg.stream_token(stream_resp)
+    await msg.send()
+
+
+# -------------------------
+# FastAPI (API Mode)
+# -------------------------
+app = FastAPI()
+global_pipeline = None  # Keep one pipeline in memory
+
+@app.post("/upload/")
+async def upload_file(file: UploadFile):
+    global global_pipeline
+    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file.filename.split('.')[-1]}") as tmp:
+        shutil.copyfileobj(file.file, tmp)
+        tmp_path = tmp.name
+
+    texts = process_file_path(tmp_path)
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_documents(texts, embeddings)
+
+    chat_llm = ChatOpenAI(
+        model="qwen/qwen2.5-vl-72b-instruct",
+        streaming=True,
+        temperature=0,
+        max_tokens=1024,
+        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://openrouter.ai/api/v1"),
+        openai_api_key= "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"
+    )
+
+    global_pipeline = RetrievalAugmentedQAPipeline(llm=chat_llm, vectorstore=vectorstore)
+    return JSONResponse({"status": "File uploaded and processed ✅", "filename": file.filename})
+
+@app.post("/upload_url/")
+async def upload_file_url(file_url: str = Form(...)):
+    global global_pipeline
+
+    # Download file from URL
+    response = requests.get(file_url, stream=True)
+    if response.status_code != 200:
+        return JSONResponse({"error": f"Failed to download file: {response.status_code}"}, status_code=400)
+
+    filename = file_url.split("/")[-1] or "downloaded_file.pdf"
+    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{filename.split('.')[-1]}") as tmp:
+        for chunk in response.iter_content(chunk_size=8192):
+            tmp.write(chunk)
+        tmp_path = tmp.name
+
+    texts = process_file_path(tmp_path)
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    vectorstore = FAISS.from_documents(texts, embeddings)
+
+    chat_llm = ChatOpenAI(
+        model="qwen/qwen2.5-vl-72b-instruct",
+        streaming=True,
+        temperature=0,
+        max_tokens=1024,
+        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://openrouter.ai/api/v1"),
+        openai_api_key= "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"
+    )
+
+
+    global_pipeline = RetrievalAugmentedQAPipeline(llm=chat_llm, vectorstore=vectorstore)
+    return JSONResponse({"status": "File downloaded and processed ✅", "filename": filename})
+
+@app.post("/ask/")
+async def ask_question(question: str = Form(...)):
+    global global_pipeline
+    if not global_pipeline:
+        return JSONResponse({"error": "No file uploaded yet."}, status_code=400)
+
+    result = await global_pipeline.arun_pipeline(question)
+    response_text = ""
+    async for token in result["response"]:
+        response_text += token
+
+    return JSONResponse({"answer": response_text})
+
+
+# -------------------------
+# Run both Chainlit + API
+# -------------------------
+import os
+os.environ["NGROK_AUTH_TOKEN"] = "2zuN63ZzFTYUM6ABGW4C1XJHe2x_7THDuvuKcg6fJY9h9bdCH"
+# Start ngrok tunnel
+import nest_asyncio
+from pyngrok import ngrok
+import uvicorn
+
+# Allow nested event loops (needed in Colab)
+nest_asyncio.apply()
+
+# Expose port 8000
+public_url = ngrok.connect(8000)
+print("🚀 Public FastAPI URL:", public_url.public_url)
+
+# Run app
+uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=False)