Skip to content

Commit d2dbbe8

Browse files
clstaudtclstaudt
committed
pdf upload and view
1 parent 3ec0b90 commit d2dbbe8

File tree

3 files changed

+257
-45
lines changed

3 files changed

+257
-45
lines changed

‎app.py

Lines changed: 254 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
import streamlit as st
22
import ollama
3+
import uuid
4+
from datetime import datetime
5+
from streamlit_pdf_viewer import pdf_viewer
36

7+
st.set_page_config(page_title="Ollama Chatbot", layout="wide")
48
st.title("Ollama Chatbot")
59

610
# Function to get available Ollama models
@@ -29,12 +33,162 @@ def get_ollama_models():
2933
st.error(f"An unexpected error occurred while fetching Ollama models: {e}. Make sure Ollama is running.")
3034
return []
3135

32-
# Initialize chat history and selected model
33-
if "messages" not in st.session_state:
34-
st.session_state.messages = []
36+
# Initialize chat system
37+
if "chats" not in st.session_state:
38+
st.session_state.chats = {}
39+
if "current_chat_id" not in st.session_state:
40+
st.session_state.current_chat_id = None
3541
if "selected_model" not in st.session_state:
3642
st.session_state.selected_model = None
3743

44+
# Helper functions for chat management
45+
def create_new_chat():
46+
"""Create a new chat session"""
47+
chat_id = str(uuid.uuid4())
48+
st.session_state.chats[chat_id] = {
49+
"messages": [],
50+
"created_at": datetime.now(),
51+
"title": "New Document Chat",
52+
"document_uploaded": False,
53+
"document_name": None,
54+
"document_content": None
55+
}
56+
st.session_state.current_chat_id = chat_id
57+
return chat_id
58+
59+
def get_current_messages():
60+
"""Get messages for the current chat"""
61+
if st.session_state.current_chat_id and st.session_state.current_chat_id in st.session_state.chats:
62+
return st.session_state.chats[st.session_state.current_chat_id]["messages"]
63+
return []
64+
65+
def add_message_to_current_chat(role, content):
66+
"""Add a message to the current chat"""
67+
if st.session_state.current_chat_id and st.session_state.current_chat_id in st.session_state.chats:
68+
st.session_state.chats[st.session_state.current_chat_id]["messages"].append({
69+
"role": role,
70+
"content": content
71+
})
72+
# Update chat title based on first user message
73+
if role == "user" and st.session_state.chats[st.session_state.current_chat_id]["title"] == "New Document Chat":
74+
title = content[:50] + "..." if len(content) > 50 else content
75+
st.session_state.chats[st.session_state.current_chat_id]["title"] = title
76+
77+
def delete_chat(chat_id):
78+
"""Delete a chat session"""
79+
if chat_id in st.session_state.chats:
80+
del st.session_state.chats[chat_id]
81+
if st.session_state.current_chat_id == chat_id:
82+
# Switch to another chat or create new one
83+
if st.session_state.chats:
84+
st.session_state.current_chat_id = list(st.session_state.chats.keys())[0]
85+
else:
86+
st.session_state.current_chat_id = None
87+
88+
def get_chat_preview(chat_data):
89+
"""Get a preview of the last message in the chat"""
90+
messages = chat_data.get("messages", [])
91+
document_name = chat_data.get("document_name", None)
92+
93+
if not messages:
94+
if document_name:
95+
return f"📄 {document_name}"
96+
return "No document uploaded yet"
97+
98+
last_message = messages[-1]
99+
content = last_message["content"]
100+
101+
# Truncate long messages
102+
if len(content) > 60:
103+
content = content[:60] + "..."
104+
105+
# Add role prefix
106+
if last_message["role"] == "user":
107+
return f"You: {content}"
108+
else:
109+
return content
110+
111+
def format_chat_time(created_at):
112+
"""Format the chat creation time"""
113+
now = datetime.now()
114+
diff = now - created_at
115+
116+
if diff.days > 0:
117+
return created_at.strftime("%b %d")
118+
elif diff.seconds > 3600:
119+
return created_at.strftime("%H:%M")
120+
else:
121+
minutes = diff.seconds // 60
122+
if minutes == 0:
123+
return "now"
124+
return f"{minutes}m ago"
125+
126+
# Initialize click tracking
127+
if "chat_clicked" not in st.session_state:
128+
st.session_state.chat_clicked = None
129+
if "delete_clicked" not in st.session_state:
130+
st.session_state.delete_clicked = None
131+
132+
# Create first chat if none exist
133+
if not st.session_state.chats:
134+
create_new_chat()
135+
136+
# Sidebar for chat management
137+
with st.sidebar:
138+
st.header("💬 Chat History")
139+
140+
# New Document Chat button
141+
if st.button("📄 New Document Chat", use_container_width=True, type="primary"):
142+
create_new_chat()
143+
st.rerun()
144+
145+
st.markdown("---")
146+
147+
# Display chat history
148+
if st.session_state.chats:
149+
# Sort chats by creation time (newest first)
150+
sorted_chats = sorted(
151+
st.session_state.chats.items(),
152+
key=lambda x: x[1]["created_at"],
153+
reverse=True
154+
)
155+
156+
for chat_id, chat_data in sorted_chats:
157+
is_current = chat_id == st.session_state.current_chat_id
158+
preview = get_chat_preview(chat_data)
159+
time_str = format_chat_time(chat_data["created_at"])
160+
161+
# Create a clean container for each chat
162+
with st.container():
163+
col1, col2 = st.columns([5, 1])
164+
165+
with col1:
166+
# Simple, clean button with all info
167+
button_type = "primary" if is_current else "secondary"
168+
169+
# Create a clean button label
170+
button_label = f"💬 **{chat_data['title']}**\n{preview}\n*{time_str}*"
171+
172+
if st.button(
173+
button_label,
174+
key=f"chat-{chat_id}",
175+
use_container_width=True,
176+
type=button_type
177+
):
178+
st.session_state.current_chat_id = chat_id
179+
st.rerun()
180+
181+
with col2:
182+
if st.button("🗑️", key=f"del-{chat_id}", help="Delete", type="secondary"):
183+
delete_chat(chat_id)
184+
st.rerun()
185+
186+
# Add subtle spacing
187+
st.write("")
188+
else:
189+
st.info("💡 No chats yet. Click 'New Document Chat' to start your first conversation!")
190+
191+
# Main content area
38192
# Get available models
39193
available_models = get_ollama_models()
40194

@@ -54,46 +208,102 @@ def get_ollama_models():
54208
else:
55209
st.session_state.selected_model = None # Ensure no model is selected if none are available
56210

57-
# Display chat messages from history on app rerun
58-
for message in st.session_state.messages:
59-
with st.chat_message(message["role"]):
60-
st.markdown(message["content"])
211+
# Check if current chat has a document uploaded
212+
current_chat = st.session_state.chats.get(st.session_state.current_chat_id, {})
213+
document_uploaded = current_chat.get("document_uploaded", False)
61214

62-
# React to user input
63-
if prompt := st.chat_input("What is up?"):
64-
if not st.session_state.selected_model:
65-
st.warning("Please select a model from the dropdown above. If no models are listed, check Ollama.")
66-
else:
67-
# Display user message in chat message container
68-
with st.chat_message("user"):
69-
st.markdown(prompt)
70-
# Add user message to chat history
71-
st.session_state.messages.append({"role": "user", "content": prompt})
72-
73-
try:
74-
# Display assistant response with streaming
75-
with st.chat_message("assistant"):
76-
# Prepare messages for Ollama API (list of dicts)
77-
current_conversation = []
78-
for msg in st.session_state.messages:
79-
current_conversation.append({'role': msg['role'], 'content': msg['content']})
80-
81-
def generate_response():
82-
"""Generator function for streaming response"""
83-
for chunk in ollama.chat(
84-
model=st.session_state.selected_model,
85-
messages=current_conversation,
86-
stream=True
87-
):
88-
if chunk['message']['content']:
89-
yield chunk['message']['content']
90-
91-
# Stream the response
92-
assistant_response = st.write_stream(generate_response())
215+
if not document_uploaded and st.session_state.current_chat_id:
216+
# Show document upload interface
217+
st.markdown("## 📄 Upload Your Document")
218+
st.markdown("To start chatting, please upload a PDF document first.")
219+
220+
uploaded_file = st.file_uploader(
221+
"Choose a PDF file",
222+
type=['pdf'],
223+
help="Drag and drop a PDF file here or click to browse",
224+
key=f"uploader_{st.session_state.current_chat_id}"
225+
)
226+
227+
if uploaded_file is not None:
228+
# Update chat with document info
229+
st.session_state.chats[st.session_state.current_chat_id]["document_uploaded"] = True
230+
st.session_state.chats[st.session_state.current_chat_id]["document_name"] = uploaded_file.name
231+
st.session_state.chats[st.session_state.current_chat_id]["document_content"] = uploaded_file.getvalue()
232+
st.session_state.chats[st.session_state.current_chat_id]["title"] = f"📄 {uploaded_file.name}"
233+
234+
st.success(f"✅ Document '{uploaded_file.name}' uploaded successfully!")
235+
st.info("💬 You can now start asking questions about your document below.")
236+
st.rerun()
237+
238+
else:
239+
# Show current document info if uploaded
240+
if document_uploaded:
241+
with st.expander("📄 Current Document", expanded=False):
242+
# Shorten long file names for display
243+
full_name = current_chat.get('document_name', 'Unknown')
244+
display_name = full_name
245+
if len(full_name) > 40:
246+
display_name = full_name[:37] + "..."
247+
248+
st.write(f"**Document:** {display_name}")
93249

94-
# Add assistant response to chat history
95-
st.session_state.messages.append({"role": "assistant", "content": assistant_response})
96-
except ollama.ResponseError as e:
97-
st.error(f"Ollama API Error during chat: {e.error} (Status code: {e.status_code})")
98-
except Exception as e:
99-
st.error(f"Error communicating with Ollama during chat: {e}")
250+
# Display PDF viewer using streamlit-pdf-viewer with maximum size
251+
document_content = current_chat.get('document_content')
252+
if document_content:
253+
pdf_viewer(
254+
input=document_content,
255+
width="100%",
256+
height=800,
257+
render_text=True,
258+
key=f"pdf_viewer_{st.session_state.current_chat_id}"
259+
)
260+
261+
# Display chat messages from history on app rerun
262+
current_messages = get_current_messages()
263+
for message in current_messages:
264+
with st.chat_message(message["role"]):
265+
st.markdown(message["content"])
266+
267+
# React to user input - only show if document is uploaded
268+
if document_uploaded:
269+
if prompt := st.chat_input("Ask a question about your document..."):
270+
if not st.session_state.selected_model:
271+
st.warning("Please select a model from the dropdown above. If no models are listed, check Ollama.")
272+
else:
273+
# Display user message in chat message container
274+
with st.chat_message("user"):
275+
st.markdown(prompt)
276+
# Add user message to chat history
277+
add_message_to_current_chat("user", prompt)
278+
279+
try:
280+
# Display assistant response with streaming
281+
with st.chat_message("assistant"):
282+
# Prepare messages for Ollama API (list of dicts)
283+
current_conversation = []
284+
for msg in get_current_messages():
285+
current_conversation.append({'role': msg['role'], 'content': msg['content']})
286+
287+
def generate_response():
288+
"""Generator function for streaming response"""
289+
for chunk in ollama.chat(
290+
model=st.session_state.selected_model,
291+
messages=current_conversation,
292+
stream=True
293+
):
294+
if chunk['message']['content']:
295+
yield chunk['message']['content']
296+
297+
# Stream the response
298+
assistant_response = st.write_stream(generate_response())
299+
300+
# Add assistant response to chat history
301+
add_message_to_current_chat("assistant", assistant_response)
302+
except ollama.ResponseError as e:
303+
st.error(f"Ollama API Error during chat: {e.error} (Status code: {e.status_code})")
304+
except Exception as e:
305+
st.error(f"Error communicating with Ollama during chat: {e}")
306+
else:
307+
# Show a message when no document is uploaded
308+
if st.session_state.current_chat_id:
309+
st.info("📄 Please upload a PDF document above to start chatting.")

‎environment.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ dependencies:
1010
# You can add other conda packages here
1111
- pip:
1212
- ollama
13+
- streamlit-pdf-viewer
1314
- -e . # Installs the ragnarok package in editable mode
1415
# If you have packages not on conda, list them under pip:
1516
# - package_from_pip

‎requirements.txt

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
streamlit
2-
ollama
2+
ollama
3+
streamlit-pdf-viewer

0 commit comments

Comments
 (0)