"""
Gradio UI for Context Thread Agent - Enterprise Edition
Professional document analysis with killer features
"""
import gradio as gr
import json
import tempfile
import os
import html
from pathlib import Path
from typing import Tuple, List, Dict
from src.models import Cell, CellType
from datetime import datetime
from src.parser import NotebookParser
from src.dependencies import ContextThreadBuilder
from src.indexing import FAISSIndexer
from src.retrieval import RetrievalEngine, ContextBuilder
from src.reasoning import ContextualAnsweringSystem
from src.intent import ContextThreadEnricher
from src.groq_integration import GroqReasoningEngine
import pandas as pd
class NotebookAgentUI:
"""Enterprise-grade Gradio UI for the Context Thread Agent."""
def __init__(self):
self.current_thread = None
self.current_indexer = None
self.current_engine = None
self.answering_system = None
self.conversation_history = []
self.groq_client = None
self.keypoints_generated = False
self.keypoints_cache = None
self.current_file_name = None
self.data_profile = None
self.current_file_path = None
self.current_file_ext = None
# Initialize Groq client
try:
self.groq_client = GroqReasoningEngine()
except Exception as e:
print(f"Warning: Groq not initialized: {e}")
def load_notebook(self, notebook_file) -> Tuple[str, bool, str, str]:
"""Load and index a notebook or Excel file."""
try:
if notebook_file is None:
return "โ No file provided", False, "", ""
# Save uploaded file temporarily
with tempfile.NamedTemporaryFile(suffix=Path(notebook_file).suffix if isinstance(notebook_file, str) else ".ipynb", delete=False) as f:
if isinstance(notebook_file, str):
f.write(open(notebook_file, 'rb').read())
else:
f.write(notebook_file.read())
temp_path = f.name
file_ext = Path(temp_path).suffix.lower()
if file_ext == '.ipynb':
parser = NotebookParser()
result = parser.parse_file(temp_path)
cells = result['cells']
elif file_ext in ['.xlsx', '.xls']:
cells = self._excel_to_cells(temp_path)
else:
return "โ Unsupported file type. Please upload .ipynb or .xlsx/.xls", False, "", ""
# Build context thread
builder = ContextThreadBuilder(
notebook_name=Path(temp_path).stem,
thread_id=f"thread_{id(self)}"
)
builder.add_cells(cells)
self.current_thread = builder.build()
# Enrich with intents
enricher = ContextThreadEnricher(infer_intents=True)
self.current_thread = enricher.enrich(self.current_thread)
# Index
self.current_indexer = FAISSIndexer()
self.current_indexer.add_multiple(self.current_thread.units)
# Setup retrieval and reasoning
self.current_engine = RetrievalEngine(self.current_thread, self.current_indexer)
self.answering_system = ContextualAnsweringSystem(self.current_engine)
# Reset conversation
self.conversation_history = []
self.keypoints_generated = False
self.keypoints_cache = None
# Store file info for later use
self.current_file_path = temp_path
self.current_file_ext = file_ext
# Get appropriate preview based on file type
if file_ext in ['.xlsx', '.xls']:
notebook_preview = self.get_excel_display(temp_path)
else:
notebook_preview = self.get_notebook_display()
# Cleanup for non-Excel files
Path(temp_path).unlink()
status_msg = f"""
### โ
File Loaded Successfully!
**Document Statistics:**
- Total sections: {len(cells)}
- Code sections: {sum(1 for c in cells if c.cell_type == CellType.CODE)}
- Documentation: {sum(1 for c in cells if c.cell_type == CellType.MARKDOWN)}
- Indexed & Ready: โ
You can now:
- ๐ Browse the document in the viewer
- ๐ Generate key insights (recommended)
- โ Ask any questions about the content
"""
return status_msg, True, notebook_preview, ""
except Exception as e:
return f"โ Error loading file: {str(e)}", False, "", ""
def generate_keypoints(self) -> str:
"""Generate key points summary using Groq."""
if not self.answering_system:
return "โ No document loaded."
if self.keypoints_cache:
return self.keypoints_cache
try:
# Get comprehensive context
all_context = []
for unit in self.current_thread.units[:30]: # First 30 cells
all_context.append(f"### {unit.cell.cell_id} [{unit.cell.cell_type}]")
if unit.intent and unit.intent != "[Pending intent inference]":
all_context.append(f"Intent: {unit.intent}")
source_text = unit.cell.source if isinstance(unit.cell.source, str) else ''.join(unit.cell.source)
all_context.append(source_text[:500])
if unit.cell.outputs:
for output in unit.cell.outputs[:1]:
if 'text' in output:
raw_out = output['text']
if isinstance(raw_out, list):
raw_out = '\n'.join(raw_out)
all_context.append(f"Output: {raw_out[:200]}")
all_context.append("---")
context_text = "\n".join(all_context)
# Use Groq to generate keypoints
if self.groq_client:
result = self.groq_client.generate_keypoints(context_text, max_points=12)
if result["success"]:
self.keypoints_cache = f"## ๐ Key Insights & Summary\n\n{result['keypoints']}"
self.keypoints_generated = True
return self.keypoints_cache
else:
return f"โ {result['keypoints']}"
else:
return "โ Groq client not available. Please check your API key."
except Exception as e:
return f"โ Error generating keypoints: {str(e)}"
def set_groq_key(self, api_key: str, enable: bool) -> str:
"""Set or clear the Groq API key and reinitialize the Groq client at runtime."""
try:
if not enable:
# Disable Groq usage
self.groq_client = None
os.environ.pop("GROQ_API_KEY", None)
return "โ
Groq disabled. The system will use fallback reasoning."
if not api_key or api_key.strip() == "":
return "โ Please provide a valid Groq API key to enable Groq."
# Try to initialize Groq with the provided key
self.groq_client = GroqReasoningEngine(api_key=api_key.strip())
os.environ["GROQ_API_KEY"] = api_key.strip()
return "โ
Groq enabled successfully. Using Groq for reasoning."
except Exception as e:
self.groq_client = None
return f"โ Could not initialize Groq: {str(e)}"
def get_notebook_display(self) -> str:
"""Get Google Colab-like styled notebook content."""
if not self.current_thread:
return "No document loaded."
display = """
"""
code_cells = sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.CODE)
markdown_cells = sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.MARKDOWN)
cells_with_output = sum(1 for u in self.current_thread.units if u.cell.outputs)
display += f"""
{len(self.current_thread.units)}
Total Cells
{markdown_cells}
Documentation
{cells_with_output}
With Output
"""
for i, unit in enumerate(self.current_thread.units, 1):
cell_type_str = "CODE" if unit.cell.cell_type == CellType.CODE else "MARKDOWN"
cell_type_class = "code" if unit.cell.cell_type == CellType.CODE else "markdown"
display += f"""
"""
if unit.cell.cell_type == CellType.CODE:
# Escape HTML special characters and preserve whitespace
# Handle source as either string or list
source_text = unit.cell.source if isinstance(unit.cell.source, str) else ''.join(unit.cell.source)
code = html.escape(source_text)
display += f'
\n'
else:
# Handle source as either string or list
source_text = unit.cell.source if isinstance(unit.cell.source, str) else ''.join(unit.cell.source)
display += f'
{source_text}
\n'
if unit.cell.outputs:
display += '
\n'
display += '
Output
\n'
for output in unit.cell.outputs[:2]:
if 'text' in output:
raw_out = output['text']
if isinstance(raw_out, list):
raw_out = '\n'.join(raw_out)
output_text = html.escape(str(raw_out)[:300])
display += f'
{output_text}\n'
elif 'data' in output and 'text/plain' in output['data']:
raw_out = output['data']['text/plain']
if isinstance(raw_out, list):
raw_out = '\n'.join(raw_out)
output_text = html.escape(str(raw_out)[:300])
display += f'
{output_text}\n'
display += '
\n'
display += """
"""
display += """
๐ Document Profile & Analytics
Comprehensive analysis of your notebook
"""
# Calculate metrics
total_cells = len(self.current_thread.units)
code_cells = sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.CODE)
markdown_cells = total_cells - code_cells
cells_with_output = sum(1 for u in self.current_thread.units if u.cell.outputs)
cells_with_intent = sum(1 for u in self.current_thread.units if u.intent and u.intent != "[Pending intent inference]")
total_lines = sum(len(u.cell.source.split('\n')) for u in self.current_thread.units)
avg_cell_size = total_lines // max(code_cells, 1)
profile += f"""
๐ Key Metrics
Total Cells: {total_cells}
Code Cells: {code_cells}
Documentation: {markdown_cells}
Cells with Output: {cells_with_output}
Total Lines: {total_lines}
Avg Cell Size: {avg_cell_size} lines
๐ก Code Quality Insights
"""
# Quality analysis
insights = []
if cells_with_output / max(code_cells, 1) > 0.8:
insights.append("โ
Excellent output coverage: Most cells produce outputs")
if cells_with_intent / total_cells > 0.7:
insights.append("โ
Well-structured workflow: Clear intent in most cells")
if code_cells < markdown_cells:
insights.append("โ
Well documented: Good documentation-to-code ratio")
if total_lines > 500:
insights.append("โ ๏ธ
Large notebook: Consider breaking into smaller modules")
if avg_cell_size > 30:
insights.append("โ ๏ธ
Large cells: Some cells could be smaller for clarity")
if not insights:
insights.append("โน๏ธ Standard notebook structure detected")
for insight in insights:
profile += f"
{insight}
\n"
profile += """
๐ Intent Distribution
"""
intent_counts = {}
for unit in self.current_thread.units:
if unit.intent and unit.intent != "[Pending intent inference]":
intent = unit.intent.split()[0] # Get first word of intent
intent_counts[intent] = intent_counts.get(intent, 0) + 1
for intent, count in sorted(intent_counts.items(), key=lambda x: x[1], reverse=True):
profile += f"
โข {intent}: {count} cells
\n"
profile += """
๐ฆ Dependencies & Imports
"""
imports = set()
for unit in self.current_thread.units:
if unit.cell.cell_type == CellType.CODE:
source = unit.cell.source if isinstance(unit.cell.source, str) else ''.join(unit.cell.source)
if 'import ' in source:
for line in source.split('\n'):
if line.strip().startswith(('import ', 'from ')):
# Extract module name
module = line.split('import')[0].replace('from', '').strip()
if module:
imports.add(module)
if imports:
for imp in sorted(imports)[:10]:
profile += f"
โข {imp}
\n"
else:
profile += "
No imports detected
\n"
profile += """
"""
return profile
def export_analysis(self) -> str:
"""Export analysis results."""
if not self.current_thread:
return "โ No document loaded."
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"analysis_{self.current_file_name or 'notebook'}_{timestamp}.md"
# Create markdown report
report = f"""# Document Analysis Report
Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
## Executive Summary
{self.keypoints_cache or "Key insights would be generated here."}
## Key Metrics
- Total Cells: {len(self.current_thread.units)}
- Code Cells: {sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.CODE)}
- Documentation Cells: {sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.MARKDOWN)}
## Questions Asked
"""
for msg in self.conversation_history:
if msg["role"] == "user":
report += f"\n- {msg['content'][:100]}"
# Save to file
with open(filename, 'w') as f:
f.write(report)
return f"โ
Report exported to `{filename}`"
def advanced_search(self, search_term: str) -> str:
"""Advanced search across all cells."""
if not self.current_thread or not search_term:
return "โ No document loaded or search term empty."
results = []
search_lower = search_term.lower()
for i, unit in enumerate(self.current_thread.units, 1):
source_text = unit.cell.source if isinstance(unit.cell.source, str) else ''.join(unit.cell.source)
if search_lower in source_text.lower():
results.append({
"cell": i,
"type": unit.cell.cell_type,
"intent": unit.intent,
"snippet": source_text[:150]
})
if not results:
return f"No results found for '{search_term}'"
output = f"๐ Found {len(results)} matches for '{search_term}'
\n"
for r in results[:10]:
output += f"""
Cell {r['cell']} [{r['type'].upper()}] {r['intent']}
{r['snippet']}...
"""
return output
def get_recommendations(self) -> str:
"""Generate smart recommendations."""
if not self.current_thread:
return "โ No document loaded."
recommendations = """
โญ AI-Powered Recommendations
"""
recs = []
code_cells = sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.CODE)
markdown_cells = sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.MARKDOWN)
if code_cells > 20:
recs.append("๐ Consider modularizing code into separate files/functions")
if markdown_cells == 0:
recs.append("๐ Add documentation cells for better clarity")
if len(self.current_thread.units) > 50:
recs.append("๐ This notebook is large - consider splitting into multiple notebooks")
# Check for common issues
large_cells = sum(1 for u in self.current_thread.units if len(u.cell.source) > 1000)
if large_cells > 0:
recs.append(f"โ๏ธ {large_cells} cells are very large - consider breaking them down")
cells_without_output = sum(1 for u in self.current_thread.units if u.cell.cell_type == CellType.CODE and not u.cell.outputs)
if cells_without_output > code_cells * 0.3:
recs.append("โ ๏ธ Many code cells don't have outputs - ensure cells are executable")
if not recs:
recs.append("โ
Notebook follows best practices!")
for i, rec in enumerate(recs, 1):
recommendations += f'{i}. {rec}
\n'
return recommendations
def _excel_to_cells(self, excel_path: str) -> List[Cell]:
"""Convert Excel file to notebook-like cells."""
from src.models import Cell, CellType
cells = []
xl = pd.ExcelFile(excel_path)
# Add overview cell
cells.append(Cell(
cell_id="excel_overview",
cell_type=CellType.MARKDOWN,
source=f"# Excel Document Analysis\n\nSheets: {', '.join(xl.sheet_names)}\nTotal Sheets: {len(xl.sheet_names)}",
outputs=[]
))
for sheet_name in xl.sheet_names:
df = xl.parse(sheet_name)
# Sheet header
cells.append(Cell(
cell_id=f"sheet_{sheet_name}_header",
cell_type=CellType.MARKDOWN,
source=f"## Sheet: {sheet_name}\n\n**Dimensions:** {df.shape[0]} rows ร {df.shape[1]} columns",
outputs=[]
))
# Column info
col_info = "\n".join([f"- {col}: {dtype}" for col, dtype in df.dtypes.items()])
cells.append(Cell(
cell_id=f"sheet_{sheet_name}_columns",
cell_type=CellType.MARKDOWN,
source=f"### Columns\n{col_info}",
outputs=[]
))
# Data preview
cells.append(Cell(
cell_id=f"data_{sheet_name}_preview",
cell_type=CellType.CODE,
source=f"# Preview of {sheet_name}\ndf_{sheet_name}.head(10)",
outputs=[{"data": {"text/plain": df.head(10).to_string()}}]
))
# Statistics
if df.select_dtypes(include=['number']).shape[1] > 0:
stats = df.describe().to_string()
cells.append(Cell(
cell_id=f"stats_{sheet_name}",
cell_type=CellType.CODE,
source=f"# Statistics for {sheet_name}\ndf_{sheet_name}.describe()",
outputs=[{"data": {"text/plain": stats}}]
))
return cells
def get_excel_display(self, excel_path: str) -> str:
"""Get Microsoft Excel-like styled spreadsheet content."""
xl = pd.ExcelFile(excel_path)
sheet_names = xl.sheet_names
if not sheet_names:
return "No sheets found in Excel file."
primary_sheet = sheet_names[0]
df = xl.parse(primary_sheet)
display = """
"""
display += f"""
{len(df.columns)}
Columns
{df.memory_usage(deep=True).sum() / 1024:.1f} KB
Size
{df.isnull().sum().sum()}
Missing
๐ Data Summary: {len(df)} rows ร {len(df.columns)} columns | Dtypes: {', '.join(map(str, df.dtypes.unique()))}
"""
for col in df.columns:
display += f" | {col} | \n"
display += """
"""
for idx, row in df.head(100).iterrows():
display += f" \n \n"
for col in df.columns:
value = row[col]
if pd.isna(value):
display += " | โ | \n"
else:
if isinstance(value, (int, float)):
formatted_value = f"{value:,.2f}" if isinstance(value, float) else str(value)
else:
formatted_value = str(value)[:50]
display += f" {formatted_value} | \n"
display += "
\n"
if len(df) > 100:
display += f"""
|
... and {len(df) - 100} more rows
|
"""
display += """
๐งต Context Thread Agent
AI-Powered Document Analysis & Q&A System
""")
with gr.Row():
with gr.Column(scale=2):
gr.Markdown("""
## ๐ฏ What is Context Thread Agent?
Context Thread Agent is an **intelligent document analysis platform** that helps you understand and extract insights from complex Jupyter notebooks and Excel spreadsheets. Using advanced AI (powered by **Groq LLM**), it provides:
### ๐ Major Use Cases:
- **๐ Data Analysis Review**: Understand complex analytical workflows instantly
- **๐ Code Audit**: Verify assumptions and logic in data science notebooks
- **๐ Excel Report Analysis**: Extract insights from large spreadsheets
- **๐ค Automated Documentation**: Generate summaries and key findings
- **๐ก Knowledge Extraction**: Ask questions about methodology and results
- **๐ Dependency Tracking**: Understand how different parts connect
- **โ
Quality Assurance**: Validate calculations and transformations
### โจ Key Features:
- โ **100% Grounded Answers** - No hallucinations, only facts from your document
- โ **Citation-Based** - Every answer references specific cells
- โ **Context-Aware** - Understands relationships between code sections
- โ **Conversation Memory** - Maintains context across questions
- โ **Key Insights Generation** - AI-powered summary of main points
- โ **Fast & Free** - Powered by Groq's lightning-fast inference
""")
with gr.Column(scale=1):
gr.HTML("""
๐ค Quick Start
Upload your document and start exploring
""")
file_input = gr.File(
label="Upload Your Document",
file_types=[".ipynb", ".xlsx", ".xls"],
type="filepath",
elem_classes="upload-input"
)
upload_btn = gr.Button(
"๐ค Upload & Analyze",
variant="primary",
size="lg",
scale=2
)
upload_status = gr.Markdown("### ๐ Status\n\nReady to upload...")
# Groq status - show only status if enabled, otherwise show input
if agent.groq_client:
groq_status = gr.Markdown("### ๐ Groq Configuration\n\nโ
**Groq is enabled and ready!**\n\nYour Groq API key has been loaded from environment. Advanced reasoning will be used for analysis.")
# Hidden inputs for compatibility
groq_key_input = gr.Textbox(visible=False)
groq_toggle = gr.Checkbox(visible=False)
set_groq_btn = gr.Button(visible=False)
else:
# Show input if Groq not enabled
groq_key_input = gr.Textbox(
label="Groq API Key",
placeholder="Paste your Groq key (gsk_...)",
type="password"
)
groq_toggle = gr.Checkbox(label="Use Groq for reasoning", value=False)
set_groq_btn = gr.Button("Set Groq Key", variant="secondary")
groq_status = gr.Markdown("โ ๏ธ **Groq not configured.** Add your key and click 'Set Groq Key' to enable advanced reasoning.")
# Wire the set key button only if inputs are visible
set_groq_btn.click(agent.set_groq_key, inputs=[groq_key_input, groq_toggle], outputs=[groq_status])
gr.Markdown("---")
# Main interface (hidden until upload)
with gr.Column(visible=False) as main_interface:
gr.Markdown("## ๐ผ Analysis Workspace")
with gr.Row():
# Left side: Document viewer
with gr.Column(scale=1):
gr.Markdown("### ๐ Document Viewer")
with gr.Tabs():
with gr.Tab("๐ Content"):
notebook_display = gr.HTML(
value="",
label="Document Content",
elem_classes="notebook-viewer"
)
with gr.Tab("๐ Key Points"):
keypoints_btn = gr.Button(
"๐ Generate Key Insights",
variant="secondary",
size="lg"
)
gr.Markdown("*This may take 10-30 seconds for comprehensive analysis...*")
keypoints_display = gr.Markdown(
value="",
label="Key Insights"
)
with gr.Tab("๐ Analytics"):
analytics_btn = gr.Button("๐ Generate Profile", variant="secondary", size="lg")
analytics_display = gr.Markdown(value="", label="Analytics")
with gr.Tab("โญ Recommendations"):
rec_btn = gr.Button("๐ก Get Recommendations", variant="secondary", size="lg")
rec_display = gr.Markdown(value="", label="Recommendations")
with gr.Tab("๐ Advanced Search"):
search_input = gr.Textbox(
label="Search Term",
placeholder="Search in all cells...",
lines=1
)
search_btn = gr.Button("๐ Search", variant="secondary")
search_display = gr.Markdown(value="", label="Search Results")
with gr.Tab("๐ฅ Export"):
export_btn = gr.Button("๐ฅ Export Analysis Report", variant="secondary", size="lg")
export_display = gr.Markdown(value="", label="Export Status")
# Right side: Q&A Interface
with gr.Column(scale=1):
gr.Markdown("### ๐ฌ Ask Questions")
chatbot = gr.Chatbot(
label="Conversation",
height=500,
elem_classes="chat-box"
)
with gr.Row():
query_input = gr.Textbox(
label="Your Question",
placeholder="e.g., 'What are the main findings?' or 'Why was Q4 data removed?'",
lines=2,
scale=4
)
ask_btn = gr.Button("๐ค Ask", variant="primary", scale=1)
gr.Markdown("""
**๐ก Example Questions:**
- What is this document about?
- What are the key findings?
- Why was [specific data] removed?
- How was [metric] calculated?
- What patterns were found?
- Are there any data quality issues?
""")
# Event handlers
def on_upload(file):
status, show_interface, notebook_content, keypoints = agent.load_notebook(file)
return (
status,
gr.update(visible=show_interface),
notebook_content,
keypoints
)
upload_btn.click(
fn=on_upload,
inputs=[file_input],
outputs=[upload_status, main_interface, notebook_display, keypoints_display]
)
# Keypoints generation with loading state
def generate_with_loading():
return "โณ **Analyzing document and generating insights...**\n\nThis may take 10-30 seconds depending on document complexity."
keypoints_btn.click(
fn=generate_with_loading,
inputs=[],
outputs=[keypoints_display]
).then(
fn=agent.generate_keypoints,
inputs=[],
outputs=[keypoints_display]
)
# Analytics tab
analytics_btn.click(
fn=agent.generate_data_profile,
inputs=[],
outputs=[analytics_display]
)
# Recommendations tab
rec_btn.click(
fn=agent.get_recommendations,
inputs=[],
outputs=[rec_display]
)
# Advanced search
search_btn.click(
fn=agent.advanced_search,
inputs=[search_input],
outputs=[search_display]
)
# Export
export_btn.click(
fn=agent.export_analysis,
inputs=[],
outputs=[export_display]
)
# Q&A interaction
ask_btn.click(
fn=agent.ask_question,
inputs=[query_input, chatbot],
outputs=[chatbot, query_input]
)
query_input.submit(
fn=agent.ask_question,
inputs=[query_input, chatbot],
outputs=[chatbot, query_input]
)
return demo
if __name__ == "__main__":
demo = create_gradio_app()
demo.launch(
server_name="0.0.0.0",
server_port=7860,
share=True
)