Supratim0406/AI-Code-Assistant

1import streamlit as st2from langchain_client import LangChainClient3import tempfile, os4from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader5 6st.set_page_config(page_title="AI Code Assistant", layout="wide")7 8# --- Session state ---9if "messages" not in st.session_state:10    st.session_state.messages = []11if "mode" not in st.session_state:12    st.session_state.mode = None13if "uploaded_texts" not in st.session_state:14    st.session_state.uploaded_texts = {}  # {filename: extracted_text}15 16st.title("🤖 AI Code Assistant (Gemini + LangChain + Streamlit)")17 18# --- Sidebar Upload ---19st.sidebar.header("📂 Project Files")20uploaded_files = st.sidebar.file_uploader(21    "Upload project files",22    accept_multiple_files=True,23    type=[24        "py", "js", "ts", "java", "cpp", "c", "html", "css", "json",25        "go", "rb", "php", "cs", "txt", "md", "docx", "pdf"26    ]27)28 29# --- Sidebar Mode Dropdown ---30st.sidebar.header("⚙️ Mode Selection")31mode = st.sidebar.selectbox(32    "Choose assistant mode (required)",33    [34        "General", "Code Analysis", "Code Generator", "Debugger",35        "Code Guide", "Optimization", "Explain Code",36        "Project Builder", "Documentation"37    ],38    index=039)40st.session_state.mode = mode41 42# Init LangChain client43LC = LangChainClient(mode=st.session_state.mode)44st.caption(f"🟢 Assistant is running in **{mode}** mode")45 46# --- File Content Extractor (with LangChain loaders) ---47def extract_text(file):48    suffix = os.path.splitext(file.name)[-1].lower()49    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:50        tmp.write(file.read())51        tmp_path = tmp.name52 53    try:54        if suffix == ".docx":55            loader = Docx2txtLoader(tmp_path)56        elif suffix == ".pdf":57            loader = PyPDFLoader(tmp_path)58        else:59            # plain text or code file60            with open(tmp_path, "r", encoding="utf-8", errors="ignore") as f:61                return f.read()62 63        docs = loader.load()64        return "\n".join([doc.page_content for doc in docs])65    finally:66        os.unlink(tmp_path)  # cleanup temp file67 68# --- Save uploaded file text in session (no auto-analysis) ---69if uploaded_files:70    for f in uploaded_files:71        if f.name not in st.session_state.uploaded_texts:72            try:73                st.session_state.uploaded_texts[f.name] = extract_text(f)74            except Exception as e:75                st.session_state.uploaded_texts[f.name] = f"⚠ Could not read {f.name}: {e}"76    st.sidebar.success(f"{len(st.session_state.uploaded_texts)} file(s) ready for analysis.")77 78# --- Analyze Button ---79if uploaded_files:80    if st.sidebar.button("🔍 Analyze Uploaded Files"):81        combined_texts = "\n\n".join(82            [f"📄 {name}:\n{text[:3000]}" for name, text in st.session_state.uploaded_texts.items()]83        )84        # Send to LLM without polluting chat history85        with st.chat_message("assistant"):86            with st.spinner("Analyzing uploaded files..."):87                reply = LC.chat(88                    st.session_state.messages + [{"role": "user", "content": combined_texts}]89                )90                st.markdown(reply)91        st.session_state.messages.append({"role": "assistant", "content": reply})92 93# --- Chat history ---94for m in st.session_state.messages:95    with st.chat_message(m["role"]):96        st.markdown(m["content"])97 98# --- Chat Input ---99if prompt := st.chat_input("Ask me to generate code, debug, review, or explain..."):100    st.session_state.messages.append({"role": "user", "content": prompt})101    with st.chat_message("user"):102        st.markdown(prompt)103 104    with st.chat_message("assistant"):105        with st.spinner("Thinking..."):106            reply = LC.chat(st.session_state.messages)107            st.markdown(reply)108 109    st.session_state.messages.append({"role": "assistant", "content": reply})