sswetha02/Notebookclone

1import streamlit as st2import os3import logging4from typing import List, Dict, Any5from pathlib import Path6import re7from io import BytesIO8import base649import pandas as pd10from PyPDF2 import PdfReader11from docx import Document12from openai import OpenAI13import requests14import moviepy.editor as mp15from pptx import Presentation16from pptx.util import Inches17 18logging.basicConfig(level=logging.INFO)19logger = logging.getLogger(__name__)20 21# Load environment22from dotenv import load_dotenv23load_dotenv()24 25# OpenAI Setup26OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")27if not OPENAI_API_KEY:28    st.error("Missing OPENAI_API_KEY in .env or secrets.")29    st.stop()30client = OpenAI(api_key=OPENAI_API_KEY)31 32# Languages for Audio (Natural OpenAI TTS)33LANGUAGES = {34    "English": "nova",35    "Hindi": "nova",36    "Bengali": "nova",37    "Tamil": "nova",38    "Telugu": "nova",39    "Kannada": "nova",40    "Marathi": "nova",41    "Gujarati": "nova",42    "Malayalam": "nova",43    "Punjabi": "nova",44}45 46# Simplified Interactive Citations (Finance-Focused)47def create_interactive_citations(response_text: str, sources_used: List[Dict[str, Any]]) -> str:48    citation_map = {}49    for idx, source in enumerate(sources_used, 1):50        citation_map[str(idx)] = source51 52    def replace_citation(match):53        num = match.group(1)54        if num in citation_map:55            source = citation_map[num]56            chunk_content = source.get('content', "Content not available")[:300] + "..." if len(source.get('content', "")) > 300 else source.get('content', "")57            source_info = f"Source: {source.get('source_file', 'Unknown')}"58            if source.get('page_number'):59                source_info += f", Page: {source['page_number']}"60            chunk_content_escaped = chunk_content.replace('<', '&lt;').replace('>', '&gt;').replace('\n', '<br>').replace('"', '&quot;')61            source_info_escaped = source_info.replace('<', '&lt;').replace('>', '&gt;').replace('"', '&quot;')62            return f'''<span class="citation-number">63                {num}64                <div class="citation-tooltip">65                    <div class="tooltip-source">{source_info_escaped}</div>66                    <div class="tooltip-content">{chunk_content_escaped}</div>67                </div>68            </span>'''69        return match.group(0)70 71    return re.sub(r'\[(\d+)\]', replace_citation, response_text)72 73# Simplified Processing (No Vector DB, Direct Context)74def process_uploaded_files(uploaded_files):75    sources = []76    for uploaded_file in uploaded_files:77        try:78            content = ""79            if uploaded_file.type == "application/pdf":80                reader = PdfReader(uploaded_file)81                for idx, page in enumerate(reader.pages, 1):82                    text = page.extract_text() or ""83                    sources.append({84                        'source_file': uploaded_file.name,85                        'page_number': idx,86                        'content': text87                    })88                content = "\n".join([s['content'] for s in sources])89            elif uploaded_file.type in ["application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/msword"]:90                doc = Document(uploaded_file)91                content = "\n".join([para.text for para in doc.paragraphs])92                sources.append({'source_file': uploaded_file.name, 'content': content})93            elif uploaded_file.type in ["text/csv", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"]:94                df = pd.read_csv(uploaded_file) if uploaded_file.type == "text/csv" else pd.read_excel(uploaded_file)95                content = df.to_markdown()96                sources.append({'source_file': uploaded_file.name, 'content': content})97            elif uploaded_file.type == "text/plain":98                content = uploaded_file.read().decode("utf-8")99                sources.append({'source_file': uploaded_file.name, 'content': content})100            st.success(f"Processed {uploaded_file.name}")101            return sources102        except Exception as e:103            st.error(f"Failed to process {uploaded_file.name}: {str(e)}")104    return []105 106def process_urls(urls_text):107    urls = [url.strip() for url in urls_text.split('\n') if url.strip()]108    sources = []109    for url in urls:110        try:111            content = requests.get(url).text112            sources.append({'source_file': url, 'content': content[:15000]})113            st.success(f"Processed {url}")114        except Exception as e:115            st.error(f"Failed to process {url}: {str(e)}")116    return sources117 118# Finance-Focused Generation119def generate_response(query, sources):120    ctx = "\n\n".join([s['content'] for s in sources])[:12000]121    response = client.chat.completions.create(122        model="gpt-4o-mini",123        messages=[124            {"role": "system", "content": "You are a financial analyst. Provide insights, risks, and summaries."},125            {"role": "user", "content": f"Context:\n{ctx}\n\nQuery: {query}"}126        ]127    ).choices[0].message.content128    interactive = create_interactive_citations(response, sources)129    return {"response": response, "interactive": interactive, "sources_used": sources}130 131# Audio/Video/Slide Generation (Finance Overviews)132def generate_audio(text, language):133    voice = LANGUAGES.get(language, "nova")134    response = client.audio.speech.create(model="tts-1", voice=voice, input=text)135    return response.content136 137def generate_video(summary):138    clip = mp.TextClip(summary[:300], fontsize=28, color='black', bg_color='white', size=(1280,720))139    audio_buf = BytesIO(generate_audio(summary, "English"))140    audio_clip = mp.AudioFileClip(audio_buf)141    video = clip.set_audio(audio_clip).set_duration(audio_clip.duration)142    buf = BytesIO()143    video.write_videofile(buf, fps=24, codec="libx264", logger=None)144    buf.seek(0)145    return buf146 147def generate_slide(summary):148    prs = Presentation()149    slide = prs.slides.add_slide(prs.slide_layouts[1])150    title = slide.shapes.title151    title.text = "Financial Overview"152    tf = title.text_frame153    tf.paragraphs[0].font.size = Pt(32)154    content = slide.placeholders[1]155    tf = content.text_frame156    tf.text = summary[:500]157    buf = BytesIO()158    prs.save(buf)159    buf.seek(0)160    return buf161 162# Streamlit UI (NotebookLM-Style)163st.set_page_config(page_title="FinAgent NotebookLM", layout="wide")164 165st.markdown("""166<style>167@import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;500;700&display=swap');168.stApp { background:#f8f9fa; font-family: 'Roboto', sans-serif; color:#202124; }169.panel { background:white; border:1px solid #dadce0; border-radius:8px; padding:16px; height:calc(100vh - 150px); overflow-y:auto; }170.source-pill { background:#f1f3f4; border:1px solid #dadce0; border-radius:20px; padding:8px 16px; margin:8px 0; font-size:14px; display:flex; align-items:center; gap:8px; }171.studio-tile { background:#f8f9fa; border:1px solid #dadce0; border-radius:8px; padding:16px; height:80px; text-align:center; transition:0.2s; }172.studio-tile:hover { background:#e8f0fe; border-color:#1a73e8; }173</style>174""", unsafe_allow_html=True)175 176# Session State177if 'sources' not in st.session_state: st.session_state.sources = []178if 'chat_history' not in st.session_state: st.session_state.chat_history = []179if 'notebook_title' not in st.session_state: st.session_state.notebook_title = "Untitled notebook"180 181# Header182col_title, col_create = st.columns([4, 1])183with col_title: st.subheader(st.session_state.notebook_title)184with col_create: if st.button("+ Create notebook"): st.session_state.sources = []; st.session_state.chat_history = []; st.session_state.notebook_title = "Untitled notebook"; st.rerun()185 186# Panels187col_src, col_chat, col_studio = st.columns([1, 3, 1.5], gap="small")188 189with col_src:190    st.markdown('<div class="panel">', unsafe_allow_html=True)191    st.markdown("Sources")192    uploaded = st.file_uploader("", accept_multiple_files=True, label_visibility="collapsed")193    url = st.text_input("Web link", placeholder="https://...")194    if uploaded:195        new_sources = process_uploaded_files(uploaded)196        st.session_state.sources.extend(new_sources)197        st.rerun()198    if url and st.button("Add Link"):199        new_sources = process_urls(url)200        st.session_state.sources.extend(new_sources)201        st.rerun()202    for source in st.session_state.sources:203        st.markdown(f'<div class="source-pill">📄 {source["source_file"]}</div>', unsafe_allow_html=True)204    st.markdown('</div>', unsafe_allow_html=True)205 206with col_chat:207    st.markdown('<div class="panel">', unsafe_allow_html=True)208    for role, content in st.session_state.chat_history:209        with st.chat_message(role):210            st.write(content.get('interactive', content['content']))211    if prompt := st.chat_input("Enter your financial query"):212        result = generate_response(prompt, st.session_state.sources)213        st.session_state.chat_history.append({"role": "user", "content": prompt})214        st.session_state.chat_history.append({"role": "assistant", "content": result})215        st.rerun()216    st.markdown('</div>', unsafe_allow_html=True)217 218with col_studio:219    st.markdown('<div class="panel">', unsafe_allow_html=True)220    st.markdown("Studio")221    if st.session_state.sources:222        g1, g2 = st.columns(2)223        with g1:224            if st.button("🔊 Audio Overview", use_container_width=True):225                lang = st.selectbox("Language", list(LANGUAGES.keys()))226                summary = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role":"user","content":f"Financial podcast script:\n{ctx}"}]).choices[0].message.content227                audio = generate_audio(summary, lang)228                st.audio(audio, format="audio/mp3")229