File size: 9,655 Bytes
a33458e
 
 
 
 
 
9f0d171
 
 
 
 
 
a33458e
 
 
c7d6600
 
 
 
 
 
 
 
 
 
 
 
 
 
a33458e
 
 
 
 
 
 
 
9f0d171
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48a1a2b
 
 
 
9f0d171
 
48a1a2b
9f0d171
 
 
 
 
 
 
 
 
 
 
a33458e
 
 
 
9f0d171
 
 
a33458e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9f0d171
a33458e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9f0d171
a33458e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9f0d171
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a33458e
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
import streamlit as st
import os
import sys
import tempfile
from datetime import datetime
from typing import List, Dict, Any
import time
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Add project root to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

# Use relative imports when running as part of the app package
try:
    from app.core.agent import AssistantAgent
    from app.core.ingestion import DocumentProcessor
    from app.utils.helpers import get_document_path, format_sources, save_conversation
    from app.config import LLM_MODEL, EMBEDDING_MODEL
except ImportError:
    # Fallback to direct imports if app is not recognized as a package
    sys.path.append(os.path.abspath('.'))
    from app.core.agent import AssistantAgent
    from app.core.ingestion import DocumentProcessor
    from app.utils.helpers import get_document_path, format_sources, save_conversation
    from app.config import LLM_MODEL, EMBEDDING_MODEL

# Set page config
st.set_page_config(
    page_title="Personal AI Assistant (Hugging Face)",
    page_icon="🤗",
    layout="wide"
)

# Function to initialize the agent safely
@st.cache_resource
def get_agent():
    logger.info("Initializing AssistantAgent (should only happen once)")
    try:
        return AssistantAgent()
    except Exception as e:
        logger.error(f"Error initializing agent: {e}")
        st.error(f"Could not initialize AI assistant: {str(e)}")
        # Return a dummy agent as fallback
        class DummyAgent:
            def query(self, question):
                return {
                    "answer": "I'm having trouble starting up. Please try refreshing the page.",
                    "sources": []
                }
            def add_conversation_to_memory(self, *args, **kwargs):
                pass
        return DummyAgent()

# Function to initialize document processor safely
@st.cache_resource
def get_document_processor(_agent):
    """Initialize document processor with unhashable agent parameter.
    The leading underscore in _agent tells Streamlit not to hash this parameter.
    """
    logger.info("Initializing DocumentProcessor (should only happen once)")
    try:
        return DocumentProcessor(_agent.memory_manager)
    except Exception as e:
        logger.error(f"Error initializing document processor: {e}")
        st.error(f"Could not initialize document processor: {str(e)}")
        # Return a dummy processor as fallback
        class DummyProcessor:
            def ingest_file(self, *args, **kwargs):
                return ["dummy-id"]
            def ingest_text(self, *args, **kwargs):
                return ["dummy-id"]
        return DummyProcessor()

# Initialize session state variables
if "messages" not in st.session_state:
    st.session_state.messages = []

# Initialize agent and document processor with caching to prevent multiple instances
agent = get_agent()
document_processor = get_document_processor(agent)

# App title
st.title("🤗 Personal AI Assistant (Hugging Face)")

# Create a sidebar for uploading documents and settings
with st.sidebar:
    st.header("Upload Documents")
    uploaded_file = st.file_uploader("Choose a file", type=["pdf", "txt", "csv"])
    
    if uploaded_file is not None:
        # Create a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{uploaded_file.name.split('.')[-1]}") as tmp:
            tmp.write(uploaded_file.getvalue())
            tmp_path = tmp.name
        
        if st.button("Process Document"):
            with st.spinner("Processing document..."):
                try:
                    # Get a path to store the document
                    doc_path = get_document_path(uploaded_file.name)
                    
                    # Copy the file to the documents directory
                    with open(doc_path, "wb") as f:
                        f.write(uploaded_file.getvalue())
                    
                    # Ingest the document
                    document_processor.ingest_file(tmp_path, {"original_name": uploaded_file.name})
                    
                    # Clean up the temporary file
                    os.unlink(tmp_path)
                    
                    st.success(f"Document {uploaded_file.name} processed successfully!")
                except Exception as e:
                    st.error(f"Error processing document: {str(e)}")
    
    st.header("Raw Text Input")
    text_input = st.text_area("Enter text to add to the knowledge base")
    
    if st.button("Add Text"):
        if text_input:
            with st.spinner("Adding text to knowledge base..."):
                try:
                    # Create metadata
                    metadata = {
                        "type": "manual_input",
                        "timestamp": str(datetime.now())
                    }
                    
                    # Ingest the text
                    document_processor.ingest_text(text_input, metadata)
                    
                    st.success("Text added to knowledge base successfully!")
                except Exception as e:
                    st.error(f"Error adding text: {str(e)}")
    
    # Display model information
    st.header("Models")
    st.write(f"**LLM**: [{LLM_MODEL}](https://huggingface.co/{LLM_MODEL})")
    st.write(f"**Embeddings**: [{EMBEDDING_MODEL}](https://huggingface.co/{EMBEDDING_MODEL})")
    
    # Add Hugging Face deployment info
    st.header("Deployment")
    st.write("This app can be easily deployed to [Hugging Face Spaces](https://huggingface.co/spaces) for free hosting.")
    
    # Link to Hugging Face
    st.markdown("""
    <div style="text-align: center; margin-top: 20px;">
        <a href="https://huggingface.co" target="_blank">
            <img src="https://huggingface.co/front/assets/huggingface_logo.svg" width="200" alt="Hugging Face">
        </a>
    </div>
    """, unsafe_allow_html=True)

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])
        
        # Display sources if available
        if message["role"] == "assistant" and "sources" in message:
            with st.expander("View Sources"):
                sources = message["sources"]
                if sources:
                    for i, source in enumerate(sources, 1):
                        st.write(f"{i}. {source['file_name']}" + (f" (Page {source['page']})" if source.get('page') else ""))
                        st.text(source['content'])
                else:
                    st.write("No specific sources used.")

# Chat input
if prompt := st.chat_input("Ask a question..."):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message
    with st.chat_message("user"):
        st.write(prompt)
    
    # Generate response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                # Add retry mechanism for vector store issues
                max_retries = 3
                for attempt in range(max_retries):
                    try:
                        response = agent.query(prompt)
                        break
                    except Exception as e:
                        if "already accessed by another instance" in str(e) and attempt < max_retries - 1:
                            logger.warning(f"Vector store access conflict, retrying ({attempt+1}/{max_retries})...")
                            time.sleep(1)  # Wait before retrying
                        else:
                            raise
                
                answer = response["answer"]
                sources = response["sources"]
                
                # Display the response
                st.write(answer)
                
                # Display sources in an expander
                with st.expander("View Sources"):
                    if sources:
                        for i, source in enumerate(sources, 1):
                            st.write(f"{i}. {source['file_name']}" + (f" (Page {source['page']})" if source.get('page') else ""))
                            st.text(source['content'])
                    else:
                        st.write("No specific sources used.")
                
                # Save conversation
                save_conversation(prompt, answer, sources)
                
                # Add assistant response to chat history
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": answer,
                    "sources": sources
                })
                
                # Update the agent's memory
                agent.add_conversation_to_memory(prompt, answer)
                
            except Exception as e:
                error_msg = f"Error generating response: {str(e)}"
                logger.error(error_msg)
                st.error(error_msg)
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": "I'm sorry, I encountered an error while processing your request. Please try again or refresh the page.",
                    "sources": []
                })

# Add a footer
st.markdown("---")
st.markdown("Built with LangChain, Hugging Face, and Qdrant")

if __name__ == "__main__":
    # This is used when running the file directly
    pass