wehub-resource-sync
720 lines
34 KiB
Python
720 lines
34 KiB
Python
import streamlit as st
|
|
import os
|
|
import json
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
from src.workflows import ResearchAssistantFlow
|
|
|
|
st.set_page_config(
|
|
page_title="AI Research Assistant",
|
|
page_icon="🔬",
|
|
layout="wide",
|
|
initial_sidebar_state="expanded"
|
|
)
|
|
|
|
st.markdown("""
|
|
<style>
|
|
.main-header {
|
|
font-size: 3rem;
|
|
font-weight: bold;
|
|
text-align: center;
|
|
margin-bottom: 2rem;
|
|
background: linear-gradient(90deg, #1e3a8a, #3b82f6);
|
|
-webkit-background-clip: text;
|
|
-webkit-text-fill-color: transparent;
|
|
}
|
|
|
|
.source-card {
|
|
background: #f8fafc;
|
|
border: 1px solid #e2e8f0;
|
|
border-radius: 8px;
|
|
padding: 1rem;
|
|
margin: 0.5rem 0;
|
|
}
|
|
|
|
.citation-item {
|
|
background: #ffffff;
|
|
border-left: 4px solid #3b82f6;
|
|
padding: 0.8rem;
|
|
margin: 0.3rem 0;
|
|
border-radius: 0 4px 4px 0;
|
|
}
|
|
|
|
.status-success {
|
|
color: #059669;
|
|
font-weight: bold;
|
|
}
|
|
|
|
.status-error {
|
|
color: #dc2626;
|
|
font-weight: bold;
|
|
}
|
|
|
|
.status-warning {
|
|
color: #d97706;
|
|
font-weight: bold;
|
|
}
|
|
|
|
.metric-card {
|
|
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
|
color: white;
|
|
padding: 1rem;
|
|
border-radius: 8px;
|
|
text-align: center;
|
|
margin: 0.5rem 0;
|
|
}
|
|
</style>
|
|
""", unsafe_allow_html=True)
|
|
|
|
def initialize_session_state():
|
|
if 'assistant' not in st.session_state:
|
|
st.session_state.assistant = None
|
|
|
|
if 'chat_history' not in st.session_state:
|
|
st.session_state.chat_history = []
|
|
|
|
if 'document_processed' not in st.session_state:
|
|
st.session_state.document_processed = False
|
|
|
|
if 'processing_status' not in st.session_state:
|
|
st.session_state.processing_status = {}
|
|
|
|
if 'current_document' not in st.session_state:
|
|
st.session_state.current_document = None
|
|
|
|
if 'last_response' not in st.session_state:
|
|
st.session_state.last_response = None
|
|
|
|
def check_api_keys() -> Dict[str, bool]:
|
|
api_keys = {
|
|
'OPENAI_API_KEY': bool(os.getenv('OPENAI_API_KEY')),
|
|
'FIRECRAWL_API_KEY': bool(os.getenv('FIRECRAWL_API_KEY')),
|
|
'ZEP_API_KEY': bool(os.getenv('ZEP_API_KEY')),
|
|
'VOYAGE_API_KEY': bool(os.getenv('VOYAGE_API_KEY')),
|
|
'TENSORLAKE_API_KEY': bool(os.getenv('TENSORLAKE_API_KEY'))
|
|
}
|
|
return api_keys
|
|
|
|
|
|
class StreamlitResearchAssistant:
|
|
def __init__(self, user_id: str = "streamlit_user", thread_id: str = "streamlit_session"):
|
|
self.user_id = user_id
|
|
self.thread_id = thread_id
|
|
self.flow = None
|
|
self.initialized = False
|
|
|
|
def initialize(self) -> bool:
|
|
try:
|
|
# Initialize the flow
|
|
self.flow = ResearchAssistantFlow(
|
|
tensorlake_api_key=os.getenv("TENSORLAKE_API_KEY"),
|
|
voyage_api_key=os.getenv("VOYAGE_API_KEY"),
|
|
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
|
zep_api_key=os.getenv("ZEP_API_KEY"),
|
|
firecrawl_api_key=os.getenv("FIRECRAWL_API_KEY"),
|
|
milvus_db_path="milvus_lite.db"
|
|
)
|
|
|
|
self.initialized = True
|
|
return True
|
|
except Exception as e:
|
|
st.error(f"Failed to initialize Research Assistant: {str(e)}")
|
|
return False
|
|
|
|
def query(self, user_query: str) -> Dict[str, Any]:
|
|
if not self.initialized:
|
|
return {"error": "Research Assistant not initialized"}
|
|
|
|
try:
|
|
# Execute the flow
|
|
result = self.flow.kickoff(inputs={
|
|
"query": user_query,
|
|
"user_id": self.user_id,
|
|
"thread_id": self.thread_id
|
|
})
|
|
return result
|
|
|
|
except Exception as e:
|
|
error_msg = f"Error processing query: {e}"
|
|
return {"error": error_msg}
|
|
|
|
def create_research_assistant() -> Optional[StreamlitResearchAssistant]:
|
|
try:
|
|
assistant = StreamlitResearchAssistant()
|
|
if assistant.initialize():
|
|
return assistant
|
|
return None
|
|
except Exception as e:
|
|
st.error(f"Failed to create Research Assistant: {str(e)}")
|
|
return None
|
|
|
|
def process_uploaded_document(uploaded_file, assistant: StreamlitResearchAssistant) -> bool:
|
|
try:
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
|
|
tmp_file.write(uploaded_file.getvalue())
|
|
tmp_file.flush()
|
|
os.fsync(tmp_file.fileno())
|
|
tmp_file_path = tmp_file.name
|
|
|
|
st.session_state.processing_status = {
|
|
'stage': 'uploading',
|
|
'message': 'Uploading document...',
|
|
'progress': 0.1
|
|
}
|
|
|
|
# Process document
|
|
progress_bar = st.progress(0.1)
|
|
status_text = st.empty()
|
|
|
|
status_text.text("📄 Uploading document...")
|
|
time.sleep(0.5)
|
|
|
|
progress_bar.progress(0.3)
|
|
status_text.text("🔍 Parsing document content...")
|
|
time.sleep(1)
|
|
|
|
progress_bar.progress(0.6)
|
|
status_text.text("🧠 Generating embeddings...")
|
|
time.sleep(1)
|
|
|
|
progress_bar.progress(0.8)
|
|
status_text.text("💾 Storing in vector database...")
|
|
|
|
if assistant.initialized:
|
|
try:
|
|
results = assistant.flow.process_documents([tmp_file_path])
|
|
st.session_state.current_document = uploaded_file.name
|
|
st.session_state.document_processed = True
|
|
|
|
progress_bar.progress(1.0)
|
|
status_text.text("✅ Document processed successfully!")
|
|
|
|
os.unlink(tmp_file_path)
|
|
|
|
except Exception as e:
|
|
if os.path.exists(tmp_file_path):
|
|
os.unlink(tmp_file_path)
|
|
|
|
error_msg = str(e)
|
|
if "TensorLake" in error_msg:
|
|
raise Exception(f"Document parsing failed: {error_msg}")
|
|
elif "Embedding" in error_msg:
|
|
raise Exception(f"Embedding generation failed: {error_msg}")
|
|
elif "API" in error_msg or "key" in error_msg.lower():
|
|
raise Exception(f"API authentication failed: {error_msg}")
|
|
else:
|
|
raise Exception(f"Document processing failed: {error_msg}")
|
|
else:
|
|
if os.path.exists(tmp_file_path):
|
|
os.unlink(tmp_file_path)
|
|
raise Exception("Research Assistant not initialized")
|
|
|
|
time.sleep(1)
|
|
progress_bar.empty()
|
|
status_text.empty()
|
|
st.session_state.processing_status = {
|
|
'stage': 'completed',
|
|
'message': f'Document "{uploaded_file.name}" processed successfully',
|
|
'progress': 1.0
|
|
}
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
st.error(f"Error processing document: {str(e)}")
|
|
st.session_state.processing_status = {
|
|
'stage': 'error',
|
|
'message': f'Error: {str(e)}',
|
|
'progress': 0.0
|
|
}
|
|
return False
|
|
|
|
|
|
def display_citations_dropdown(response: Dict[str, Any], key: str):
|
|
if 'context_sources' not in response:
|
|
return
|
|
|
|
context_sources = response['context_sources']
|
|
evaluation_result = response.get('evaluation_result', {})
|
|
|
|
try:
|
|
relevant_source_keys = evaluation_result.get('relevant_sources', [])
|
|
title = "📚 **View Sources & Citations**"
|
|
|
|
with st.expander(title, expanded=False):
|
|
if 'relevant_sources' in evaluation_result:
|
|
st.markdown("#### 🎯 Source Relevance Summary")
|
|
|
|
relevant_sources = evaluation_result['relevant_sources']
|
|
relevance_scores = evaluation_result.get('relevance_scores', {})
|
|
reasoning = evaluation_result.get('reasoning', 'No reasoning provided')
|
|
|
|
col1, col2 = st.columns([1, 2])
|
|
with col1:
|
|
st.markdown("**Relevant Sources:**")
|
|
for source in relevant_sources:
|
|
score = relevance_scores.get(source, 'N/A')
|
|
if isinstance(score, (int, float)):
|
|
st.markdown(f"• **{source}**: {score:.2f}")
|
|
else:
|
|
st.markdown(f"• **{source}**: {score}")
|
|
|
|
with col2:
|
|
st.markdown("**Reasoning:**")
|
|
st.markdown(f"*{reasoning}*")
|
|
|
|
st.markdown("---")
|
|
|
|
# Only display sources that are marked as relevant by the evaluator
|
|
all_sources = [
|
|
('RAG (Documents)', context_sources.get('rag_result', {}), '📄', 'RAG'),
|
|
('Memory (History)', context_sources.get('memory_result', {}), '🧠', 'Memory'),
|
|
('Web Search', context_sources.get('web_result', {}), '🌐', 'Web'),
|
|
('ArXiv Papers', context_sources.get('tool_result', {}), '📚', 'ArXiv')
|
|
]
|
|
# Filter sources based on evaluation result
|
|
relevant_source_keys = evaluation_result.get('relevant_sources', [])
|
|
# If no evaluation result available, show all sources
|
|
if not relevant_source_keys:
|
|
sources = [(name, data, icon) for name, data, icon, key in all_sources if data]
|
|
else:
|
|
# Only show sources that were marked as relevant
|
|
sources = []
|
|
for name, data, icon, key in all_sources:
|
|
if data and key in relevant_source_keys:
|
|
sources.append((name, data, icon))
|
|
|
|
if not sources:
|
|
st.markdown("*No relevant sources found for this query.*")
|
|
return
|
|
|
|
for source_name, source_data, icon in sources:
|
|
if not source_data:
|
|
continue
|
|
|
|
if source_name == 'Memory (History)':
|
|
status = 'OK'
|
|
elif source_name == 'Web Search':
|
|
has_search_results = source_data.get('search_results')
|
|
has_explicit_status = source_data.get('status') == 'OK'
|
|
has_answer = source_data.get('answer')
|
|
has_relevance = source_data.get('relevance_assessment')
|
|
|
|
if has_search_results or has_explicit_status or (has_answer and has_relevance):
|
|
status = 'OK'
|
|
elif source_data.get('status') == 'ERROR':
|
|
status = 'ERROR'
|
|
elif source_data.get('status') == 'INSUFFICIENT_CONTEXT':
|
|
status = 'INSUFFICIENT_CONTEXT'
|
|
else:
|
|
status = 'UNKNOWN'
|
|
elif source_name == 'ArXiv Papers':
|
|
status = source_data.get('status', 'UNKNOWN')
|
|
else: # RAG
|
|
status = source_data.get('status', 'UNKNOWN')
|
|
|
|
# Create expandable section for each source
|
|
with st.expander(f"{icon} **{source_name}** ({status})", expanded=False):
|
|
|
|
if status == 'OK':
|
|
if source_name == 'Memory (History)':
|
|
context = source_data.get('context', [])
|
|
if context:
|
|
st.markdown("**Memory Context:**")
|
|
|
|
if isinstance(context, (list, tuple)):
|
|
items_to_show = context[:6]
|
|
for i, item in enumerate(items_to_show):
|
|
item_str = str(item) if item is not None else ""
|
|
if len(item_str) > 200:
|
|
truncated_item = item_str[:200] + "..."
|
|
else:
|
|
truncated_item = item_str
|
|
st.markdown(f"• {truncated_item}")
|
|
|
|
if len(context) > 6:
|
|
st.markdown(f"*...and {len(context) - 6} more items*")
|
|
else:
|
|
st.markdown(f"• {str(context)[:500]}...")
|
|
|
|
relevance = source_data.get('relevance_assessment', {})
|
|
if relevance:
|
|
citations = relevance.get('citations', [])
|
|
if citations:
|
|
st.markdown("**Citations:**")
|
|
for citation in citations:
|
|
label = citation.get('label', 'Citation')
|
|
locator = citation.get('locator', 'N/A')
|
|
st.markdown(f"• **{label}** ({locator})")
|
|
|
|
confidence = relevance.get('confidence', 'N/A')
|
|
if confidence != 'N/A':
|
|
st.markdown(f"**Confidence:** {confidence}")
|
|
|
|
elif source_name == 'Web Search':
|
|
search_results = source_data.get('search_results', [])
|
|
answer = source_data.get('answer', '')
|
|
|
|
if search_results:
|
|
st.markdown("**Web Search Results:**")
|
|
if isinstance(search_results, (list, tuple)):
|
|
results_to_show = search_results[:3]
|
|
for i, result in enumerate(results_to_show):
|
|
if isinstance(result, dict):
|
|
title = result.get('title', 'No title')
|
|
url = result.get('url', '#')
|
|
content = str(result.get('content', 'No content'))[:150]
|
|
st.markdown(f"**{i+1}. [{title}]({url})**")
|
|
st.markdown(f"*{content}...*")
|
|
st.markdown("---")
|
|
else:
|
|
st.markdown(f"**{i+1}.** {str(result)[:200]}...")
|
|
|
|
if len(search_results) > 3:
|
|
st.markdown(f"*...and {len(search_results) - 3} more results*")
|
|
else:
|
|
st.markdown(f"• {str(search_results)[:500]}...")
|
|
|
|
elif answer and answer.strip():
|
|
st.markdown("**Web Search Content:**")
|
|
if answer.startswith('**') or '**' in answer:
|
|
st.markdown(answer[:1000] + ('...' if len(answer) > 1000 else ''))
|
|
else:
|
|
st.markdown(f"```\n{answer[:500]}{'...' if len(answer) > 500 else ''}\n```")
|
|
|
|
relevance = source_data.get('relevance_assessment', {})
|
|
if relevance:
|
|
confidence = relevance.get('confidence', 'N/A')
|
|
if confidence != 'N/A':
|
|
st.markdown(f"**Confidence:** {confidence}")
|
|
|
|
citations = source_data.get('citations', [])
|
|
if citations:
|
|
st.markdown("**Citations:**")
|
|
for citation in citations:
|
|
if isinstance(citation, dict):
|
|
label = citation.get('label', 'Web Citation')
|
|
locator = citation.get('locator', '#')
|
|
if locator.startswith('http'):
|
|
st.markdown(f"• [{label}]({locator})")
|
|
else:
|
|
st.markdown(f"• **{label}** ({locator})")
|
|
else:
|
|
st.markdown(f"• {str(citation)}")
|
|
|
|
elif source_name == 'ArXiv Papers':
|
|
answer = source_data.get('answer', '')
|
|
papers = []
|
|
if answer:
|
|
try:
|
|
import json
|
|
parsed_answer = json.loads(answer)
|
|
papers = parsed_answer.get('papers', [])
|
|
except json.JSONDecodeError:
|
|
st.markdown("**ArXiv Response:**")
|
|
st.markdown(f"```\n{answer[:300]}...\n```")
|
|
|
|
if papers:
|
|
st.markdown("**Academic Papers:**")
|
|
if isinstance(papers, (list, tuple)):
|
|
papers_to_show = papers[:3]
|
|
for i, paper in enumerate(papers_to_show):
|
|
if isinstance(paper, dict):
|
|
title = paper.get('title', 'No title')
|
|
authors = paper.get('authors', [])
|
|
url = paper.get('url', '#')
|
|
abstract = str(paper.get('abstract', 'No abstract'))[:200]
|
|
|
|
st.markdown(f"**{i+1}. [{title}]({url})**")
|
|
if authors and isinstance(authors, (list, tuple)):
|
|
authors_to_show = authors[:3] if len(authors) > 3 else authors
|
|
authors_str = ', '.join(str(author) for author in authors_to_show)
|
|
if len(authors) > 3:
|
|
authors_str += f" and {len(authors) - 3} others"
|
|
st.markdown(f"*Authors: {authors_str}*")
|
|
st.markdown(f"*{abstract}...*")
|
|
st.markdown("---")
|
|
else:
|
|
st.markdown(f"**{i+1}.** {str(paper)[:200]}...")
|
|
|
|
if len(papers) > 3:
|
|
st.markdown(f"*...and {len(papers) - 3} more papers*")
|
|
else:
|
|
st.markdown(f"• {str(papers)[:500]}...")
|
|
|
|
else: # RAG or other sources
|
|
st.markdown("**Content:**")
|
|
try:
|
|
answer = source_data.get('answer', 'No answer available')
|
|
if answer is None:
|
|
st.markdown("```\nNo content available\n```")
|
|
elif isinstance(answer, (str)):
|
|
preview = answer[:300] if len(answer) > 300 else answer
|
|
ellipsis = '...' if len(answer) > 300 else ''
|
|
st.markdown(f"```\n{preview}{ellipsis}\n```")
|
|
elif isinstance(answer, (dict, list)):
|
|
try:
|
|
import json
|
|
json_str = json.dumps(answer, indent=2)
|
|
preview = json_str[:300] if len(json_str) > 300 else json_str
|
|
ellipsis = '...' if len(json_str) > 300 else ''
|
|
st.markdown(f"```json\n{preview}{ellipsis}\n```")
|
|
except Exception:
|
|
st.markdown(f"```\n{str(answer)[:300]}...\n```")
|
|
else:
|
|
answer_str = str(answer)
|
|
preview = answer_str[:300] if len(answer_str) > 300 else answer_str
|
|
ellipsis = '...' if len(answer_str) > 300 else ''
|
|
st.markdown(f"```\n{preview}{ellipsis}\n```")
|
|
|
|
except Exception as answer_error:
|
|
st.error(f"Error displaying answer: {str(answer_error)}")
|
|
st.markdown("```\nError loading content\n```")
|
|
|
|
# Show citations with enhanced metadata
|
|
citations = source_data.get('citations', [])
|
|
if citations and isinstance(citations, (list, tuple)):
|
|
st.markdown("**Citations:**")
|
|
for i, citation in enumerate(citations):
|
|
try:
|
|
if not isinstance(citation, dict):
|
|
st.markdown(f"• Citation {i+1}: {str(citation)}")
|
|
continue
|
|
|
|
label = citation.get('label', f'Citation {i+1}')
|
|
locator = citation.get('locator', 'No location')
|
|
label = str(label) if label is not None else f'Citation {i+1}'
|
|
locator = str(locator) if locator is not None else 'No location'
|
|
|
|
page_number = citation.get('page_number')
|
|
chunk_index = citation.get('chunk_index')
|
|
score = citation.get('score')
|
|
chunk_content = citation.get('content', '')
|
|
|
|
if locator.startswith('http'):
|
|
st.markdown(f"• [{label}]({locator})")
|
|
elif page_number is not None and chunk_index is not None:
|
|
score_text = f" (Score: {score:.3f})" if isinstance(score, (int, float)) else ""
|
|
st.markdown(f"**📄 Page {page_number}, Chunk {chunk_index}**{score_text}")
|
|
|
|
if chunk_content:
|
|
content_preview = chunk_content[:300] if len(chunk_content) > 300 else chunk_content
|
|
ellipsis = '...' if len(chunk_content) > 300 else ''
|
|
st.markdown(f"```\n{content_preview}{ellipsis}\n```")
|
|
else:
|
|
st.markdown("*No content preview available*")
|
|
elif 'chunk_' in locator:
|
|
st.markdown(f"• **{label}** (Document chunk)")
|
|
else:
|
|
st.markdown(f"• **{label}**")
|
|
except Exception as citation_error:
|
|
st.markdown(f"• Citation {i+1}: Error displaying citation ({str(citation_error)})")
|
|
elif citations:
|
|
st.markdown("**Citations:**")
|
|
st.markdown(f"• Raw citation data: {str(citations)[:200]}...")
|
|
|
|
# Show additional metadata
|
|
if 'retrieval_metadata' in source_data:
|
|
metadata = source_data['retrieval_metadata']
|
|
if 'retrieved_chunks' in metadata:
|
|
st.markdown(f"**Retrieved Chunks:** {metadata['retrieved_chunks']}")
|
|
if 'document_count' in metadata:
|
|
st.markdown(f"**Documents Searched:** {metadata['document_count']}")
|
|
|
|
confidence = source_data.get('confidence', 'N/A')
|
|
if confidence != 'N/A':
|
|
if isinstance(confidence, (int, float)):
|
|
st.markdown(f"**Confidence:** {confidence:.2f}")
|
|
else:
|
|
st.markdown(f"**Confidence:** {confidence}")
|
|
|
|
elif status == 'INSUFFICIENT_CONTEXT':
|
|
st.warning(f"{source_data.get('answer', 'No relevant information found')}")
|
|
|
|
else:
|
|
error_msg = source_data.get('error', source_data.get('message', source_data.get('answer', 'Unknown error')))
|
|
st.error(f"{error_msg}")
|
|
|
|
except Exception as e:
|
|
st.error(f"❌ Error displaying citations: {str(e)}")
|
|
st.caption(f"Debug info: Error type: {type(e).__name__}")
|
|
|
|
# Show raw data for debugging
|
|
with st.expander("🔍 Debug Information", expanded=False):
|
|
st.json({
|
|
"context_sources_keys": list(context_sources.keys()) if isinstance(context_sources, dict) else str(type(context_sources)),
|
|
"evaluation_result_keys": list(evaluation_result.keys()) if isinstance(evaluation_result, dict) else str(type(evaluation_result)),
|
|
"error_details": str(e)
|
|
})
|
|
|
|
|
|
def display_sidebar_document_processing():
|
|
with st.sidebar:
|
|
st.markdown("## 📄 Document Processing")
|
|
if not st.session_state.assistant:
|
|
if st.button("🚀 Initialize Research Assistant", type="primary"):
|
|
with st.spinner("Initializing..."):
|
|
assistant = create_research_assistant()
|
|
if assistant:
|
|
st.session_state.assistant = assistant
|
|
st.success("✅ Assistant initialized!")
|
|
st.rerun()
|
|
else:
|
|
st.error("❌ Failed to initialize!")
|
|
st.markdown("---")
|
|
return
|
|
|
|
# Document upload
|
|
uploaded_file = st.file_uploader(
|
|
"Upload PDF Document",
|
|
type=['pdf'],
|
|
help="Upload a PDF document to analyze"
|
|
)
|
|
|
|
if uploaded_file is not None:
|
|
col1, col2 = st.columns([3, 1])
|
|
with col1:
|
|
st.info(f"📄 **{uploaded_file.name}**")
|
|
st.caption(f"Size: {uploaded_file.size:,} bytes")
|
|
|
|
with col2:
|
|
if st.button("Process", type="primary", key="process_doc"):
|
|
with st.spinner("Processing..."):
|
|
success = process_uploaded_document(uploaded_file, st.session_state.assistant)
|
|
if success:
|
|
st.session_state.document_processed = True
|
|
st.session_state.current_document = uploaded_file.name
|
|
st.success("✅ Processed!")
|
|
st.rerun()
|
|
else:
|
|
st.error("❌ Failed!")
|
|
|
|
if st.session_state.document_processed:
|
|
st.success("✅ Document Ready")
|
|
if st.session_state.current_document:
|
|
st.caption(f"Current: {st.session_state.current_document}")
|
|
else:
|
|
st.info("📋 No document processed")
|
|
|
|
st.markdown("---")
|
|
|
|
if st.session_state.assistant and st.session_state.assistant.initialized:
|
|
st.success("🤖 Assistant: Online")
|
|
else:
|
|
st.error("🤖 Assistant: Offline")
|
|
|
|
def display_main_chat_interface():
|
|
col1, col2 = st.columns([4, 1])
|
|
|
|
with col1:
|
|
st.markdown("## 💬 Research Chat")
|
|
with col2:
|
|
if st.button("🔄 Reset Chat", type="secondary", key="reset_chat"):
|
|
st.session_state.chat_history = []
|
|
st.session_state.last_response = None
|
|
st.success("Chat reset!")
|
|
st.rerun()
|
|
|
|
if not st.session_state.document_processed:
|
|
st.warning("⚠️ Please process a document first using the sidebar.")
|
|
return
|
|
|
|
# Display chat history
|
|
for i, (query, response) in enumerate(st.session_state.chat_history):
|
|
with st.container():
|
|
# User message
|
|
st.markdown(f"**🧑 You:** {query}")
|
|
# Assistant response
|
|
if isinstance(response, dict) and 'final_response' in response:
|
|
st.markdown(f"**🤖 Assistant:** {response['final_response']}")
|
|
# Add citations dropdown
|
|
display_citations_dropdown(response, f"citations_{i}")
|
|
else:
|
|
st.markdown(f"**🤖 Assistant:** {response}")
|
|
|
|
st.markdown("---")
|
|
|
|
query = st.chat_input("Ask me anything about your document...")
|
|
|
|
if query:
|
|
# Add user message to history
|
|
with st.spinner("🔍 Researching your question..."):
|
|
try:
|
|
# Show progress steps
|
|
progress_container = st.container()
|
|
with progress_container:
|
|
st.info("📄 **Step 1/4:** Analyzing document...")
|
|
time.sleep(0.5)
|
|
st.info("🧠 **Step 2/4:** Retrieving memories...")
|
|
time.sleep(0.5)
|
|
st.info("🌐 **Step 3/4:** Searching web...")
|
|
time.sleep(0.5)
|
|
st.info("📚 **Step 4/4:** Searching academic papers...")
|
|
time.sleep(0.5)
|
|
|
|
result = st.session_state.assistant.query(query)
|
|
|
|
progress_container.empty()
|
|
# Add to chat history
|
|
st.session_state.chat_history.append((query, result))
|
|
st.session_state.last_response = result
|
|
st.rerun()
|
|
|
|
except Exception as e:
|
|
st.error(f"Error processing query: {str(e)}")
|
|
|
|
def display_initialization_message():
|
|
st.info("⚠️ Please initialize the Research Assistant using the sidebar to begin.")
|
|
|
|
def main():
|
|
initialize_session_state()
|
|
|
|
st.markdown('''
|
|
<div style="text-align: center; margin-bottom: 30px;">
|
|
<h1 style='color: #ffffff; font-size: 3rem; font-weight: bold; margin-bottom: 10px;'>
|
|
🔬 AI Research Assistant
|
|
</h1>
|
|
<div style="display: flex; justify-content: center; align-items: center; gap: 8px; margin-bottom: 20px;">
|
|
<span style='color: #64748b; font-size: 16px; font-weight: 500;'>Powered by</span>
|
|
<div style="display: flex; align-items: center; gap: 25px; margin-left: 15px;">
|
|
<a href="https://www.tensorlake.ai/" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://i.ibb.co/PZD1qrPg/tensorlake-logo.png"
|
|
alt="Tensorlake" style="height: 36px;">
|
|
</a>
|
|
<a href="https://www.getzep.com/" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://i.ibb.co/DgtgNLVQ/zep-logo.png"
|
|
alt="Zep" style="height: 32px;">
|
|
</a>
|
|
<a href="https://www.firecrawl.dev/" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://i.ibb.co/67jyMHfy/firecrawl-light-wordmark.png"
|
|
alt="Firecrawl" style="height: 28px;">
|
|
</a>
|
|
<a href="https://www.crewai.com/" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://i.ibb.co/JwmNZhCx/crewai-logo.png"
|
|
alt="CrewAI" style="height: 28px;">
|
|
</a>
|
|
<a href="https://milvus.io/" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://milvus.io/images/layout/milvus-logo.svg"
|
|
alt="Milvus" style="height: 28px;">
|
|
</a>
|
|
</div>
|
|
</div>
|
|
<p style='color: #64748b; font-size: 14px; margin-top: 10px;'>
|
|
<b>Context Engineering Workflow</b> with RAG, Web Search, Memory & Academic Research
|
|
</p>
|
|
</div>
|
|
''', unsafe_allow_html=True)
|
|
|
|
display_sidebar_document_processing()
|
|
|
|
if st.session_state.assistant and st.session_state.assistant.initialized:
|
|
display_main_chat_interface()
|
|
else:
|
|
display_initialization_message()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|