chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:37:47 +08:00
commit 7653f56fed
1422 changed files with 359026 additions and 0 deletions
@@ -0,0 +1,5 @@
TENSORLAKE_API_KEY=your_tensorlake_api_key
VOYAGE_API_KEY=your_voyage_api_key
ZEP_API_KEY=your_zep_api_key
OPENAI_API_KEY=your_openai_api_key
FIRECRAWL_API_KEY=your_firecrawl_api_key
@@ -0,0 +1 @@
3.13
+215
View File
@@ -0,0 +1,215 @@
# Context Engineering Pipeline for Research Assistant
A comprehensive research assistant that combines multiple AI agents using CrewAI Flows to provide intelligent, multi-source responses to research queries.
## Architecture Overview
This research assistant uses a multi-agent CrewAI Flow architecture with the following components:
### Core Components
1. **Document Processing & RAG Pipeline**
- TensorLake for complex document parsing with structured extraction
- Voyage Context 3 embeddings for contextualized semantic understanding
- Milvus vector database for efficient similarity search
- OpenAI GPT models with structured output formatting
2. **Memory Layer**
- Zep Cloud for persistent conversation memory
- User preference tracking
- Conversation summarization and context retrieval
3. **Web Search Integration**
- Firecrawl for real-time web search capabilities
- Retrieval of recent information not available in documents
4. **Multi-Agent Flow Architecture**
- **RAG Agent**: Searches through parsed research documents
- **Memory Agent**: Retrieves conversation history and user preferences
- **Web Search Agent**: Finds recent web-based information
- **Tool Calling Agent**: Interfaces with external APIs (extensible)
- **Evaluator Agent**: Filters and ranks context relevance
- **Synthesizer Agent**: Creates coherent final responses
## Flow Process
```mermaid
graph TD
A["User Query"] --> B["ResearchAssistantFlow<br/>Entry Point"]
B --> C["Parallel Agent Execution"]
C --> D["RAG Agent<br/>📚 Document Search"]
C --> E["Memory Agent<br/>🧠 Context Retrieval"]
C --> F["Web Search Agent<br/>🌐 Real-time Info"]
C --> G["Tool Calling Agent<br/>🔧 External APIs"]
D --> H["Context Collection<br/>📊 Aggregate Results"]
E --> H
F --> H
G --> H
H --> I["Evaluator Agent<br/>🎯 Filter Relevance"]
I --> J["Synthesizer Agent<br/>✍️ Generate Response"]
J --> K["Final Response<br/>📝 Coherent Answer"]
subgraph "RAG Pipeline Components"
D1["TensorLake<br/>Document Parser"] --> D2["Voyage Context 3<br/>Embeddings"]
D2 --> D3["Milvus Vector DB<br/>Similarity Search"]
D3 --> D
end
subgraph "Memory Components"
E1["Zep Cloud<br/>Conversation History"] --> E
E2["User Preferences<br/>Context Summaries"] --> E
end
subgraph "Generation"
J1["OpenAI GPT<br/>Structured Output"] --> J
J2["Citation Management<br/>Confidence Scoring"] --> J
end
style A fill:#e1f5fe
style K fill:#e8f5e8
style I fill:#fff3e0
style J fill:#f3e5f5
```
## Project Structure
```
context-engineering-workflow/
├── 📁 src/ # Main source code directory
│ ├── 📁 workflows/ # 🎯 Complete workflow orchestration
│ │ ├── 📄 agents.py # Agent creation from YAML configs
│ │ ├── 📄 tasks.py # Task creation from YAML configs
│ │ ├── 📄 flow.py # Main ResearchAssistantFlow
│ ├── 📁 tools/ # 🔧 All specialized tools
│ │ ├── 📄 rag_tool.py # RAG search functionality
│ │ ├── 📄 memory_tool.py # Memory retrieval tool
│ │ ├── 📄 arxiv_tool.py # ArXiv academic search
│ │ ├── 📄 web_search_tool.py # Web search via Firecrawl
│ ├── 📁 rag/ # 📚 RAG pipeline components
│ │ ├── 📄 rag_pipeline.py # Unified RAG orchestration
│ │ ├── 📄 retriever.py # Milvus vector database
│ │ ├── 📄 embeddings.py # Contextualized embeddings
│ ├── 📁 document_processing/ # 📄 Document parsing & processing
│ │ ├── 📄 doc_parser.py # TensorLake document parser
│ ├── 📁 memory/ # 🧠 Memory management
│ │ ├── 📄 memory.py # Zep memory layer
│ ├── 📁 generation/ # ✍️ Response generation
│ │ ├── 📄 generation.py # Structured response generation
│ ├── 📁 config/ # ⚙️ Configuration management
│ │ ├── 📄 config_loader.py # YAML configuration loader
├── 📁 config/ # 📋 YAML configuration files
│ ├── 📁 agents/ # Agent configurations
│ │ └── 📄 research_agents.yaml # Agent roles, goals, backstories
│ └── 📁 tasks/ # Task configurations
│ └── 📄 research_tasks.yaml # Task descriptions, expected outputs
├── 📁 data/ # 📊 Research documents (PDFs)
├── 📁 outputs/ # 📤 Generated outputs and results
├── 📄 app.py # 🌐 Streamlit web interface
├── 📄 pyproject.toml # 🔧 Project configuration
├── 📄 uv.lock # 🔒 Dependency lock file
└── 📄 README.md
```
## Installation & Setup
1. **Install dependencies:**
First, install `uv` and set up the environment:
```bash
# MacOS/Linux
curl -LsSf https://astral.sh/uv/install.sh | sh
# Windows
powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
```
Install dependencies:
```bash
# Create a new directory for our project
uv init research-assistant
cd research-assistant
# Create virtual environment and activate it
uv venv
source .venv/bin/activate # MacOS/Linux
.venv\Scripts\activate # Windows
# Install dependencies
uv sync
```
2. **Set up environment variables:**
Create a `.env` file with your API keys:
```env
TENSORLAKE_API_KEY=your_tensorlake_key
VOYAGE_API_KEY=your_voyage_key
OPENAI_API_KEY=your_openai_key
ZEP_API_KEY=your_zep_key
FIRECRAWL_API_KEY=your_firecrawl_key
```
Get the API keys here:
- [Tensorlake →](https://tensorlake.ai/)
- [Zep AI →](https://www.getzep.com/)
- [Firecrawl →](https://www.firecrawl.dev/)
- [OpenAI →](https://openai.com)
- [Voyage →](https://dashboard.voyageai.com/)
4. **Prepare documents:**
Place your research documents in the `data/` directory (PDF format supported)
## Usage
```python
uv run app.py or streamlit run app.py
```
## Key Features
### 1. Extended citations support
Each response includes comprehensive source attribution with a:
#### 🎯 Source Relevance Summary
- **Relevant Sources**: List of sources used
- **Relevance Scores**: Confidence scores (0-1) for each source
- **Reasoning**: Explanation of source selection
### 2. Multi-Source Intelligence
- Combines document knowledge, conversation memory, web search, and external APIs
- Each source operates independently and in parallel for efficiency
### 3. Intelligent Context Evaluation
- Evaluator agent filters irrelevant information
- Only relevant context is used for final response generation
### 4. Coherent Response Synthesis
- Synthesizer agent creates well-structured responses
- Proper citation and confidence scoring
- Handles insufficient context gracefully
### 5. Persistent Memory
- Conversation history stored in Zep Cloud
- User preferences and context maintained across sessions
- Agentic memory with graph-based internal representations
## API Requirements
- **TensorLake**: Document parsing and structured extraction
- **Voyage AI**: Contextualized embeddings
- **OpenAI**: Response generation with structured outputs
- **Zep Cloud**: Persistent memory and conversation management
- **Firecrawl**: Web search capabilities
## 📬 Stay Updated with Our Newsletter!
**Get a FREE Data Science eBook** 📖 with 150+ essential lessons in Data Science when you subscribe to our newsletter! Stay in the loop with the latest tutorials, insights, and exclusive resources. [Subscribe now!](https://join.dailydoseofds.com)
[![Daily Dose of Data Science Newsletter](https://github.com/patchy631/ai-engineering/blob/main/resources/join_ddods.png)](https://join.dailydoseofds.com)
---
## Contribution
Contributions are welcome! Please fork the repository and submit a pull request with your improvements.
+719
View File
@@ -0,0 +1,719 @@
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()
@@ -0,0 +1,71 @@
rag_agent:
role: >
Document Research Specialist
goal: >
Load, process, and retrieve information from research documents
backstory: >
You are an expert at managing and searching through complex research documents.
You can load documents into the system and process them to create
contextualized embeddings, and then perform semantic search to find the most relevant
information for any research query. When searching, you can provide document_paths
to load documents.
verbose: true
memory_agent:
role: >
Memory & Context Specialist
goal: >
Retrieve relevant information from conversation history and user preferences
backstory: >
You are a memory specialist who maintains context across conversations.
You can access previous discussions, user preferences, and conversation summaries
to provide relevant background context for current queries.
verbose: true
web_search_agent:
role: >
Web Research Specialist
goal: >
Search the web for recent and relevant information not available in documents
backstory: >
You are a web research expert who can find the latest information
on any topic using advanced search techniques. You specialize in finding recent developments,
news, and information that might not be available in the parsed documents.
verbose: true
arxiv_agent:
role: >
Academic Research Specialist
goal: >
Search and analyze academic papers from ArXiv to provide comprehensive research insights
backstory: >
You are an expert academic researcher with deep knowledge of scientific literature.
You can search ArXiv for relevant papers, filter by categories and authors, and provide comprehensive
analysis of academic research. You excel at finding related work, identifying research trends,
and connecting findings across different papers. You can search by title, author, abstract, category,
or perform general searches across all fields.
verbose: true
evaluator_agent:
role: >
Context Evaluation Specialist
goal: >
Evaluate and filter context from multiple sources for relevance to the query
backstory: >
You are an expert at evaluating the relevance and quality of
information from multiple sources. You can assess which pieces of context
are most relevant to answering a specific query and filter out irrelevant
information to ensure only relevant context is retrieved from each source.
verbose: true
synthesizer_agent:
role: >
Response Synthesis Specialist
goal: >
Synthesize information from multiple sources into coherent, comprehensive responses
backstory: >
You are a master at combining information from various sources
into clear, coherent, and comprehensive responses. You can weave together
insights from documents, memory, web search, and external tools to create
well-structured answers that fully address the user's query.
verbose: true
@@ -0,0 +1,76 @@
rag_search_task:
description: >
Search for information relevant to: {query}. Use query='{query}'. Search through documents that are already loaded in the vector database.
expected_output: >
JSON formatted response with status, answer, citations, and confidence score
agent: rag_agent
memory_retrieval_task:
description: >
Retrieve relevant conversation history and context for: {query}
expected_output: >
JSON formatted response with memory context and relevance assessment
agent: memory_agent
web_search_task:
description: >
Search the web for recent information and developments related to: {query}
expected_output: >
JSON formatted response with web search results and relevance assessment
agent: web_search_agent
arxiv_search_task:
description: >
Search ArXiv for academic papers related to: {query}. Find relevant research papers, authors, and recent developments in the field.
expected_output: >
JSON formatted response with ArXiv search results including paper titles, authors, abstracts, and publication details
agent: arxiv_agent
context_evaluation_task:
description: >
Evaluate the relevance of the following context sources for the query: "{query}"
Context Sources:
1. RAG Result: {rag_result}
2. Memory Result: {memory_result}
3. Web Search Result: {web_result}
4. Tool Result (ArXiv): {tool_result}
For each source, determine:
- Is it relevant to answering the query? (yes/no)
- Confidence score of relevance (0-1)
- Key information that should be included in the final response
- What information should be filtered out as irrelevant
Note: If a source has ERROR status, it should not be included in relevant_sources, but mention it in the reasoning.
Return only the relevant context that should be used for generating the final response.
IMPORTANT: Your response must strictly follow the provided JSON schema structure.
expected_output: >
JSON response with:
- relevant_sources: list of source names that are relevant (e.g., ['RAG', 'Web', 'Memory', 'ArXiv'])
- filtered_context: dictionary with only relevant information from each source
- relevance_scores: confidence scores (0-1) for each source's relevance
- reasoning: brief explanation of filtering decisions
The response must be valid JSON that conforms to the ContextEvaluationResult schema.
agent: evaluator_agent
synthesis_task:
description: >
Create a comprehensive, coherent response to the query: "{query}"
Use the following filtered and relevant context:
{filtered_context}
Guidelines:
- Synthesize information from multiple sources into a coherent narrative
- Cite sources appropriately
- Maintain accuracy and don't add information not present in the context
- Structure the response clearly and logically
expected_output: >
Well-structured response with:
- Clear answer to the user's query
- Proper citations from all relevant sources
agent: synthesizer_agent
@@ -0,0 +1,475 @@
## CHUNK NUMBER 0
## Page 1
# Attention Is All You Need
| Ashish Vaswani* Google Brain | Noam Shazeer* Google Brain | Niki Parmar* Google Research | Jakob Uszkoreit* Google Research |
| --- | --- | --- | --- |
| avaswani@google.com | noam@google.com | nikip@google.com | usz@google.com |
| Llion Jones* Google Research | Aidan N. Gomez* * University of Toronto | Łukasz Kaiser* Google Brain |
| --- | --- | --- |
| llion@google.com | aidan@cs.toronto.edu | lukaszkaiser@google.com |
Illia Polosukhin* * illia.polosukhin@gmail.com
## CHUNK NUMBER 1
## Page 1
## Abstract
The dominant sequence transduction models are based on complex recurrent or convolutional neural networks that include an encoder and a decoder. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on two machine translation tasks show these models to be superior in quality while being more parallelizable and requiring significantly less time to train. Our model achieves 28.4 BLEU on the WMT 2014 English- to-German translation task, improving over the existing best results, including ensembles, by over 2 BLEU. On the WMT 2014 English-to-French translation task, our model establishes a new single-model state-of-the-art BLEU score of 41.0 after training for 3.5 days on eight GPUs, a small fraction of the training costs of the best models from the literature.
## CHUNK NUMBER 2
## Page 2
## 1 Introduction
Recurrent neural networks, long short-term memory [12] and gated recurrent [7] neural networks in particular, have been firmly established as state of the art approaches in sequence modeling and transduction problems such as language modeling and machine translation [29, 2, 5]. Numerous efforts have since continued to push the boundaries of recurrent language models and encoder-decoder architectures [31, 21, 13].
*Equal contribution. Listing order is random. Jakob proposed replacing RNNs with self-attention and started the effort to evaluate this idea. Ashish, with Illia, designed and implemented the first Transformer models and has been crucially involved in every aspect of this work. Noam proposed scaled dot-product attention, multi-head attention and the parameter-free position representation and became the other person involved in nearly every detail. Niki designed, implemented, tuned and evaluated countless model variants in our original codebase and tensor2tensor. Llion also experimented with novel model variants, was responsible for our initial codebase, and efficient inference and visualizations. Lukasz and Aidan spent countless long days designing various parts of and implementing tensor2tensor, replacing our earlier codebase, greatly improving results and massively accelerating our research.
* Work performed while at Google Brain.
\#Work performed while at Google Research.
31st Conference on Neural Information Processing Systems (NIPS 2017), Long Beach, CA, USA.
Recurrent models typically factor computation along the symbol positions of the input and output sequences. Aligning the positions to steps in computation time, they generate a sequence of hidden states ht, as a function of the previous hidden state ht-1 and the input for position t. This inherently sequential nature precludes parallelization within training examples, which becomes critical at longer sequence lengths, as memory constraints limit batching across examples. Recent work has achieved significant improvements in computational efficiency through factorization tricks [18] and conditional computation [26], while also improving model performance in case of the latter. The fundamental constraint of sequential computation, however, remains.
Attention mechanisms have become an integral part of compelling sequence modeling and transduc- tion models in various tasks, allowing modeling of dependencies without regard to their distance in the input or output sequences [2, 16]. In all but a few cases [22], however, such attention mechanisms are used in conjunction with a recurrent network.
In this work we propose the Transformer, a model architecture eschewing recurrence and instead relying entirely on an attention mechanism to draw global dependencies between input and output. The Transformer allows for significantly more parallelization and can reach a new state of the art in translation quality after being trained for as little as twelve hours on eight P100 GPUs.
## CHUNK NUMBER 3
## Page 2
## 2 Background
The goal of reducing sequential computation also forms the foundation of the Extended Neural GPU [20], ByteNet [15] and ConvS2S [8], all of which use convolutional neural networks as basic building block, computing hidden representations in parallel for all input and output positions. In these models, the number of operations required to relate signals from two arbitrary input or output positions grows in the distance between positions, linearly for ConvS2S and logarithmically for ByteNet. This makes it more difficult to learn dependencies between distant positions [11]. In the Transformer this is reduced to a constant number of operations, albeit at the cost of reduced effective resolution due to averaging attention-weighted positions, an effect we counteract with Multi-Head Attention as described in section 3.2.
Self-attention, sometimes called intra-attention is an attention mechanism relating different positions of a single sequence in order to compute a representation of the sequence. Self-attention has been used successfully in a variety of tasks including reading comprehension, abstractive summarization, textual entailment and learning task-independent sentence representations [4, 22, 23, 19].
End-to-end memory networks are based on a recurrent attention mechanism instead of sequence- aligned recurrence and have been shown to perform well on simple-language question answering and language modeling tasks [28].
To the best of our knowledge, however, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output without using sequence- aligned RNNs or convolution. In the following sections, we will describe the Transformer, motivate self-attention and discuss its advantages over models such as [14, 15] and [8].
## CHUNK NUMBER 4
## Page 2
## 3 Model Architecture
Most competitive neural sequence transduction models have an encoder-decoder structure [5, 2, 29]. Here, the encoder maps an input sequence of symbol representations (x1, ... , In) to a sequence of continuous representations z = (z1, ... , Zn). Given z, the decoder then generates an output sequence (y1, ... , ym) of symbols one element at a time. At each step the model is auto-regressive [9], consuming the previously generated symbols as additional input when generating the next.
The Transformer follows this overall architecture using stacked self-attention and point-wise, fully connected layers for both the encoder and decoder, shown in the left and right halves of Figure 1, respectively.
## CHUNK NUMBER 5
## Page 3
### 3.1 Encoder and Decoder Stacks
Encoder: The encoder is composed of a stack of N = 6 identical layers. Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position-
2
Figure 1: The Transformer - model architecture.
### Figure
Output Probabilities
Softmax
Linear
Add & Norm
Feed Forward
Add & Norm
Add & Norm
Feed Forward
Multi-Head Attention
Nx
Nx
Add & Norm
Add & Norm
Multi-Head Attention
Masked Multi-Head Attention
Positional Encoding
Positional Encoding
Input Embedding
Output Embedding
Inputs
Outputs (shifted right)
wise fully connected feed-forward network. We employ a residual connection [10] around each of the two sub-layers, followed by layer normalization [1]. That is, the output of each sub-layer is LayerNorm(x+ Sublayer(x)), where Sublayer(x) is the function implemented by the sub-layer itself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension dmodel = 512.
Decoder: The decoder is also composed of a stack of N = 6 identical layers. In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head attention over the output of the encoder stack. Similar to the encoder, we employ residual connections around each of the sub-layers, followed by layer normalization. We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i.
## CHUNK NUMBER 6
## Page 3
### 3.2 Attention
An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.
## CHUNK NUMBER 7
## Page 4
#### 3.2.1 Scaled Dot-Product Attention
We call our particular attention "Scaled Dot-Product Attention" (Figure 2). The input consists of queries and keys of dimension dk, and values of dimension dy. We compute the dot products of the
3
Scaled Dot-Product Attention
### Figure
1
MatMul
1
SoftMax
1
Mask (opt.)
1
Scale
1
MatMul
1
1
Q
K
V
Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several attention layers running in parallel.
### Figure
Multi-Head Attention
Linear
Concat
Scaled Dot-Product Attention
h
Linear
Linear
Linear
V
K
Q
query with all keys, divide each by Vdk, and apply a softmax function to obtain the weights on the values.
In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix Q. The keys and values are also packed together into matrices K and V. We compute the matrix of outputs as:
Attention(Q, K, V) = softmax(
Vdk QKT V (1)
The two most commonly used attention functions are additive attention [2], and dot-product (multi- plicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of . Additive attention computes the compatibility function using a feed-forward network with Vdk a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.
While for small values of dk the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of dk [3]. We suspect that for large values of dk, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients 4. To counteract this effect, we scale the dot products by.
## CHUNK NUMBER 8
## Page 5
#### 3.2.2 Multi-Head Attention
Instead of performing a single attention function with dmodel-dimensional keys, values and queries, we found it beneficial to linearly project the queries, keys and values h times with different, learned linear projections to dk, dk and dy dimensions, respectively. On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding dy-dimensional output values. These are concatenated and once again projected, resulting in the final values, as depicted in Figure 2.
Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.
4To illustrate why the dot products get large, assume that the components of q and k are independent random variables with mean 0 and variance 1. Then their dot product, q . k = >iz1 qiki, has mean 0 and variance dk -
4
MultiHead(Q,K,V) = Concat(head1, ... , headh)Wº where headi = Attention(Qwº, KWK, VW)
Where the projections are parameter matrices Wo E Rdmodel X dk , W.K € Rdmodel Xdk , WY € Rdmodel X dy and WO E Rhdy Xdmodel
In this work we employ h = 8 parallel attention layers, or heads. For each of these we use dk = dy = dmodel/h = 64. Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.
## CHUNK NUMBER 9
## Page 5
#### 3.2.3 Applications of Attention in our Model
The Transformer uses multi-head attention in three different ways:
· In "encoder-decoder attention" layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This allows every position in the decoder to attend over all positions in the input sequence. This mimics the typical encoder-decoder attention mechanisms in sequence-to-sequence models such as [31, 2, 8].
· The encoder contains self-attention layers. In a self-attention layer all of the keys, values and queries come from the same place, in this case, the output of the previous layer in the encoder. Each position in the encoder can attend to all positions in the previous layer of the encoder.
· Similarly, self-attention layers in the decoder allow each position in the decoder to attend to all positions in the decoder up to and including that position. We need to prevent leftward information flow in the decoder to preserve the auto-regressive property. We implement this inside of scaled dot-product attention by masking out (setting to -oo) all values in the input of the softmax which correspond to illegal connections. See Figure 2.
## CHUNK NUMBER 10
## Page 5
### 3.3 Position-wise Feed-Forward Networks
In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.
FFN(x)=max(0,xW1+b1)W2+b2 (2)
While the linear transformations are the same across different positions, they use different parameters from layer to layer. Another way of describing this is as two convolutions with kernel size 1. The dimensionality of input and output is dmodel = 512, and the inner-layer has dimensionality dff =2048.
## CHUNK NUMBER 11
## Page 5
### 3.4 Embeddings and Softmax
Similarly to other sequence transduction models, we use learned embeddings to convert the input tokens and output tokens to vectors of dimension dmodel. We also use the usual learned linear transfor- mation and softmax function to convert the decoder output to predicted next-token probabilities. In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to [24]. In the embedding layers, we multiply those weights by Vdmodel.
## CHUNK NUMBER 12
## Page 6
### 3.5 Positional Encoding
Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence. To this end, we add "positional encodings" to the input embeddings at the
5
Table 1: Maximum path lengths, per-layer complexity and minimum number of sequential operations for different layer types. n is the sequence length, d is the representation dimension, k is the kernel size of convolutions and r the size of the neighborhood in restricted self-attention.
| Layer Type | Complexity per Layer | Sequential Operations | Maximum Path Length |
| --- | --- | --- | --- |
| Self-Attention | O(n2 · d) | O(1) | O(1) |
| Recurrent | O(n · d2) | O(n) | O(n) |
| Convolutional | O(k · n · d2) | O(1) | O(logk (n)) |
| Self-Attention (restricted) | O(r . n . d) | O(1) | O(n/r) |
bottoms of the encoder and decoder stacks. The positional encodings have the same dimension dmodel as the embeddings, so that the two can be summed. There are many choices of positional encodings, learned and fixed [8].
In this work, we use sine and cosine functions of different frequencies:
PE(pos,2i) = sin(pos/100002i/dmodel ) PE(pos,2i+1) = cos(pos/100002i/dmodel )
where pos is the position and i is the dimension. That is, each dimension of the positional encoding corresponds to a sinusoid. The wavelengths form a geometric progression from 27 to 10000 . 27. We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset k, PEpos+k can be represented as a linear function of PE, pos·
We also experimented with using learned positional embeddings [8] instead, and found that the two versions produced nearly identical results (see Table 3 row (E)). We chose the sinusoidal version because it may allow the model to extrapolate to sequence lengths longer than the ones encountered during training.
## CHUNK NUMBER 13
## Page 7
## 4 Why Self-Attention
In this section we compare various aspects of self-attention layers to the recurrent and convolu- tional layers commonly used for mapping one variable-length sequence of symbol representations (x1, ... , In) to another sequence of equal length (z1, ... , Zn), with Ti, Zi E Rd, such as a hidden layer in a typical sequence transduction encoder or decoder. Motivating our use of self-attention we consider three desiderata.
One is the total computational complexity per layer. Another is the amount of computation that can be parallelized, as measured by the minimum number of sequential operations required.
The third is the path length between long-range dependencies in the network. Learning long-range dependencies is a key challenge in many sequence transduction tasks. One key factor affecting the ability to learn such dependencies is the length of the paths forward and backward signals have to traverse in the network. The shorter these paths between any combination of positions in the input and output sequences, the easier it is to learn long-range dependencies [11]. Hence we also compare the maximum path length between any two input and output positions in networks composed of the different layer types.
As noted in Table 1, a self-attention layer connects all positions with a constant number of sequentially executed operations, whereas a recurrent layer requires O(n) sequential operations. In terms of computational complexity, self-attention layers are faster than recurrent layers when the sequence length n is smaller than the representation dimensionality d, which is most often the case with sentence representations used by state-of-the-art models in machine translations, such as word-piece [31] and byte-pair [25] representations. To improve computational performance for tasks involving very long sequences, self-attention could be restricted to considering only a neighborhood of size r in
6
the input sequence centered around the respective output position. This would increase the maximum path length to O(n/r). We plan to investigate this approach further in future work.
A single convolutional layer with kernel width k < n does not connect all pairs of input and output positions. Doing so requires a stack of O(n/k) convolutional layers in the case of contiguous kernels, or O(logk(n)) in the case of dilated convolutions [15], increasing the length of the longest paths between any two positions in the network. Convolutional layers are generally more expensive than recurrent layers, by a factor of k. Separable convolutions [6], however, decrease the complexity considerably, to O(k . n . d + n . d2). Even with k = n, however, the complexity of a separable convolution is equal to the combination of a self-attention layer and a point-wise feed-forward layer, the approach we take in our model.
As side benefit, self-attention could yield more interpretable models. We inspect attention distributions from our models and present and discuss examples in the appendix. Not only do individual attention heads clearly learn to perform different tasks, many appear to exhibit behavior related to the syntactic and semantic structure of the sentences.
## CHUNK NUMBER 14
## Page 7
## 5 Training
This section describes the training regime for our models.
## CHUNK NUMBER 15
## Page 7
### 5.1 Training Data and Batching
We trained on the standard WMT 2014 English-German dataset consisting of about 4.5 million sentence pairs. Sentences were encoded using byte-pair encoding [3], which has a shared source- target vocabulary of about 37000 tokens. For English-French, we used the significantly larger WMT 2014 English-French dataset consisting of 36M sentences and split tokens into a 32000 word-piece vocabulary [31]. Sentence pairs were batched together by approximate sequence length. Each training batch contained a set of sentence pairs containing approximately 25000 source tokens and 25000 target tokens.
## CHUNK NUMBER 16
## Page 7
### 5.2 Hardware and Schedule
We trained our models on one machine with 8 NVIDIA P100 GPUs. For our base models using the hyperparameters described throughout the paper, each training step took about 0.4 seconds. We trained the base models for a total of 100,000 steps or 12 hours. For our big models,(described on the bottom line of table 3), step time was 1.0 seconds. The big models were trained for 300,000 steps (3.5 days).
## CHUNK NUMBER 17
## Page 7
### 5.3 Optimizer
We used the Adam optimizer [17] with 31 = 0.9, 2 = 0.98 and € = 10-9. We varied the learning rate over the course of training, according to the formula:
lrate = dmodel . min(step_num-0.5, step_num . warmup_steps-1.5) (3)
This corresponds to increasing the learning rate linearly for the first warmup_steps training steps, and decreasing it thereafter proportionally to the inverse square root of the step number. We used warmup_steps = 4000.
## CHUNK NUMBER 18
## Page 8
### 5.4 Regularization
We employ three types of regularization during training:
Residual Dropout We apply dropout [27] to the output of each sub-layer, before it is added to the sub-layer input and normalized. In addition, we apply dropout to the sums of the embeddings and the positional encodings in both the encoder and decoder stacks. For the base model, we use a rate of Pdrop =0.1.
7
Table 2: The Transformer achieves better BLEU scores than previous state-of-the-art models on the English-to-German and English-to-French newstest2014 tests at a fraction of the training cost.
| Model | BLEU | | Training Cost (FLOPs) | |
| --- | --- | --- | --- | --- |
| | EN-DE | EN-FR | EN-DE | EN-FR |
| ByteNet [15] | 23.75 | | | |
| Deep-Att + PosUnk [32] | | 39.2 | | 1.0 . 1020 |
| GNMT + RL [31] | 24.6 | 39.92 | 2.3 . 1019 | 1.4 . 1020 |
| ConvS2S [8] | 25.16 | 40.46 | 9.6 . 1018 | 1.5 . 1020 |
| MoE [26] | 26.03 | 40.56 | 2.0 · 1019 | 1.2 . 1020 |
| Deep-Att + PosUnk Ensemble [32] | | 40.4 | | 8.0 . 1020 |
| GNMT + RL Ensemble [31] | 26.30 | 41.16 | 1.8 . 1020 | 1.1 . 1021 |
| ConvS2S Ensemble [8] | 26.36 | 41.29 | 7.7 . 1019 | 1.2 . 1021 |
| Transformer (base model) | 27.3 | 38.1 | 3.3 . 1018 | |
| Transformer (big) | 28.4 | 41.0 | 2.3 . 1019 | |
Label Smoothing During training, we employed label smoothing of value Els = 0.1 [30]. This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.
## CHUNK NUMBER 19
## Page 8
## 6 Results
### 6.1 Machine Translation
On the WMT 2014 English-to-German translation task, the big transformer model (Transformer (big) in Table 2) outperforms the best previously reported models (including ensembles) by more than 2.0 BLEU, establishing a new state-of-the-art BLEU score of 28.4. The configuration of this model is listed in the bottom line of Table 3. Training took 3.5 days on 8 P100 GPUs. Even our base model surpasses all previously published models and ensembles, at a fraction of the training cost of any of the competitive models.
On the WMT 2014 English-to-French translation task, our big model achieves a BLEU score of 41.0, outperforming all of the previously published single models, at less than 1/4 the training cost of the previous state-of-the-art model. The Transformer (big) model trained for English-to-French used dropout rate Pdrop = 0.1, instead of 0.3.
For the base models, we used a single model obtained by averaging the last 5 checkpoints, which were written at 10-minute intervals. For the big models, we averaged the last 20 checkpoints. We used beam search with a beam size of 4 and length penalty & = 0.6 [31]. These hyperparameters were chosen after experimentation on the development set. We set the maximum output length during inference to input length + 50, but terminate early when possible [31].
Table 2 summarizes our results and compares our translation quality and training costs to other model architectures from the literature. We estimate the number of floating point operations used to train a model by multiplying the training time, the number of GPUs used, and an estimate of the sustained single-precision floating-point capacity of each GPU 5.
## CHUNK NUMBER 20
## Page 9
### 6.2 Model Variations
To evaluate the importance of different components of the Transformer, we varied our base model in different ways, measuring the change in performance on English-to-German translation on the development set, newstest2013. We used beam search as described in the previous section, but no checkpoint averaging. We present these results in Table 3.
In Table 3 rows (A), we vary the number of attention heads and the attention key and value dimensions, keeping the amount of computation constant, as described in Section 3.2.2. While single-head attention is 0.9 BLEU worse than the best setting, quality also drops off with too many heads.
5We used values of 2.8, 3.7, 6.0 and 9.5 TFLOPS for K80, K40, M40 and P100, respectively.
8
Table 3: Variations on the Transformer architecture. Unlisted values are identical to those of the base model. All metrics are on the English-to-German translation development set, newstest2013. Listed perplexities are per-wordpiece, according to our byte-pair encoding, and should not be compared to per-word perplexities.
| | N | dmodel | dff | h | dk | dv | Pdrop | Els | train steps | PPL (dev) | BLEU (dev) | params ×106 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| base | 6 | 512 | 2048 | 8 | 64 | 64 | 0.1 | 0.1 | 100K | 4.92 | 25.8 | 65 |
| (A) | | | | 1 | 512 | 512 | | | | 5.29 | 24.9 | |
| | | | | 4 | 128 | 128 | | | | 5.00 | 25.5 | |
| | | | | 16 | 32 | 32 | | | | 4.91 | 25.8 | |
| | | | | 32 | 16 | 16 | | | | 5.01 | 25.4 | |
| (B) | | | | | 16 | | | | | 5.16 | 25.1 | 58 |
| | | | | | 32 | | | | | 5.01 | 25.4 | 60 |
| (C) | 2 | | | | | | | | | 6.11 | 23.7 | 36 |
| | 4 | | | | | | | | | 5.19 | 25.3 | 50 |
| | 8 | | | | | | | | | 4.88 | 25.5 | 80 |
| | | 256 | | | 32 | 32 | | | | 5.75 | 24.5 | 28 |
| | | 1024 | | | 128 | 128 | | | | 4.66 | 26.0 | 168 |
| | | | 1024 | | | | | | | 5.12 | 25.4 | 53 |
| | | | 4096 | | | | | | | 4.75 | 26.2 | 90 |
| (D) | | | | | | | 0.0 | | | 5.77 | 24.6 | |
| | | | | | | | 0.2 | | | 4.95 | 25.5 | |
| | | | | | | | | 0.0 | | 4.67 | 25.3 | |
| | | | | | | | | 0.2 | | 5.47 | 25.7 | |
| (E) | positional embedding | | | | | instead of | sinusoids | | | 4.92 | 25.7 | |
| big | 6 | 1024 | 4096 | 16 | | | 0.3 | | 300K | 4.33 | 26.4 | 213 |
In Table 3 rows (B), we observe that reducing the attention key size dk hurts model quality. This suggests that determining compatibility is not easy and that a more sophisticated compatibility function than dot product may be beneficial. We further observe in rows (C) and (D) that, as expected, bigger models are better, and dropout is very helpful in avoiding over-fitting. In row (E) we replace our sinusoidal positional encoding with learned positional embeddings [8], and observe nearly identical results to the base model.
## CHUNK NUMBER 21
## Page 10
## 7 Conclusion
In this work, we presented the Transformer, the first sequence transduction model based entirely on attention, replacing the recurrent layers most commonly used in encoder-decoder architectures with multi-headed self-attention.
For translation tasks, the Transformer can be trained significantly faster than architectures based on recurrent or convolutional layers. On both WMT 2014 English-to-German and WMT 2014 English-to-French translation tasks, we achieve a new state of the art. In the former task our best model outperforms even all previously reported ensembles.
We are excited about the future of attention-based models and plan to apply them to other tasks. We plan to extend the Transformer to problems involving input and output modalities other than text and to investigate local, restricted attention mechanisms to efficiently handle large inputs and outputs such as images, audio and video. Making generation less sequential is another research goals of ours.
The code we used to train and evaluate our models is available at https://github.com/ tensorflow/tensor2tensor.
Acknowledgements We are grateful to Nal Kalchbrenner and Stephan Gouws for their fruitful comments, corrections and inspiration.
9
## CHUNK NUMBER 22
## Page 11
## References
[1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607.06450, 2016.
[2] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. CoRR, abs/1409.0473, 2014.
[3] Denny Britz, Anna Goldie, Minh-Thang Luong, and Quoc V. Le. Massive exploration of neural machine translation architectures. CoRR, abs/1703.03906, 2017.
[4] Jianpeng Cheng, Li Dong, and Mirella Lapata. Long short-term memory-networks for machine reading. arXiv preprint arXiv: 1601.06733, 2016.
[5] Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. CoRR, abs/1406.1078, 2014.
[6] Francois Chollet. Xception: Deep learning with depthwise separable convolutions. arXiv preprint arXiv: 1610.02357, 2016.
[7] Junyoung Chung, Çaglar Gülçehre, Kyunghyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. CoRR, abs/1412.3555, 2014.
[8] Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N. Dauphin. Convolu- tional sequence to sequence learning. arXiv preprint arXiv: 1705.03122v2, 2017.
[9] Alex Graves. Generating sequences with recurrent neural networks. arXiv preprint arXiv: 1308.0850, 2013.
[10] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016.
[11] Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, and Jürgen Schmidhuber. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001.
[12] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735-1780, 1997.
[13] Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv: 1602.02410, 2016.
[14] Łukasz Kaiser and Ilya Sutskever. Neural GPUs learn algorithms. In International Conference on Learning Representations (ICLR), 2016.
[15] Nal Kalchbrenner, Lasse Espeholt, Karen Simonyan, Aaron van den Oord, Alex Graves, and Ko- ray Kavukcuoglu. Neural machine translation in linear time. arXiv preprint arXiv: 1610.10099v2, 2017.
[16] Yoon Kim, Carl Denton, Luong Hoang, and Alexander M. Rush. Structured attention networks. In International Conference on Learning Representations, 2017.
[17] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[18] Oleksii Kuchaiev and Boris Ginsburg. Factorization tricks for LSTM networks. arXiv preprint arXiv: 1703.10722, 2017.
[19] Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. arXiv preprint arXiv: 1703.03130, 2017.
[20] Samy Bengio Łukasz Kaiser. Can active memory replace attention? In Advances in Neural Information Processing Systems, (NIPS), 2016.
10
[21] Minh-Thang Luong, Hieu Pham, and Christopher D Manning. Effective approaches to attention- based neural machine translation. arXiv preprint arXiv: 1508.04025, 2015.
[22] Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model. In Empirical Methods in Natural Language Processing, 2016.
[23] Romain Paulus, Caiming Xiong, and Richard Socher. A deep reinforced model for abstractive summarization. arXiv preprint arXiv: 1705.04304, 2017.
[24] Ofir Press and Lior Wolf. Using the output embedding to improve language models. arXiv preprint arXiv: 1608.05859, 2016.
[25] Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint arXiv: 1508.07909, 2015.
[26] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 1701.06538, 2017.
[27] Nitish Srivastava, Geoffrey E Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdi- nov. Dropout: a simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(1):1929-1958, 2014.
[28] Sainbayar Sukhbaatar, arthur szlam, Jason Weston, and Rob Fergus. End-to-end memory networks. In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems 28, pages 2440-2448. Curran Associates, Inc., 2015.
[29] Ilya Sutskever, Oriol Vinyals, and Quoc VV Le. Sequence to sequence learning with neural networks. In Advances in Neural Information Processing Systems, pages 3104-3112, 2014.
[30] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. CoRR, abs/1512.00567, 2015.
[31] Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. Google's neural machine translation system: Bridging the gap between human and machine translation. arXiv preprint arXiv: 1609.08144, 2016.
[32] Jie Zhou, Ying Cao, Xuguang Wang, Peng Li, and Wei Xu. Deep recurrent models with fast-forward connections for neural machine translation. CoRR, abs/1606.04199, 2016.
11
@@ -0,0 +1,308 @@
[
{
"data": {
"paper": {
"abstract": "The dominant sequence transduction models are based on complex recurrent or convolutional neural networks that include an encoder and a decoder. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on two machine translation tasks show these models to be superior in quality while being more parallelizable and requiring significantly less time to train. Our model achieves 28.4 BLEU on the WMT 2014 English- to-German translation task, improving over the existing best results, including ensembles, by over 2 BLEU. On the WMT 2014 English-to-French translation task, our model establishes a new single-model state-of-the-art BLEU score of 41.0 after training for 3.5 days on eight GPUs, a small fraction of the training costs of the best models from the literature.",
"abstract_citation": [
{
"page_number": 1,
"x1": 142.0,
"x2": 469.0,
"y1": 359.0,
"y2": 500.0
}
],
"authors": [
"Ashish Vaswani",
"Noam Shazeer",
"Niki Parmar",
"Jakob Uszkoreit",
"Llion Jones",
"Aidan N. Gomez",
"\u0141ukasz Kaiser",
"Illia Polosukhin"
],
"authors_citation": [
{
"page_number": 1,
"x1": 111.0,
"x2": 222.0,
"y1": 181.0,
"y2": 208.0
},
{
"page_number": 1,
"x1": 222.0,
"x2": 315.0,
"y1": 181.0,
"y2": 207.0
},
{
"page_number": 1,
"x1": 315.0,
"x2": 413.0,
"y1": 181.0,
"y2": 207.0
},
{
"page_number": 1,
"x1": 413.0,
"x2": 501.0,
"y1": 181.0,
"y2": 207.0
},
{
"page_number": 1,
"x1": 122.0,
"x2": 222.0,
"y1": 231.0,
"y2": 257.0
},
{
"page_number": 1,
"x1": 222.0,
"x2": 351.0,
"y1": 231.0,
"y2": 257.0
},
{
"page_number": 1,
"x1": 351.0,
"x2": 490.0,
"y1": 231.0,
"y2": 257.0
},
{
"page_number": 1,
"x1": 237.0,
"x2": 374.0,
"y1": 284.0,
"y2": 307.0
}
],
"key_findings": [
"The Transformer model achieves a BLEU score of 28.4 on the WMT 2014 English-to-German translation task, outperforming previous models by over 2 BLEU.",
"On the WMT 2014 English-to-French translation task, the Transformer model achieves a BLEU score of 41.0, setting a new state-of-the-art for single models."
],
"key_findings_citation": [
{
"page_number": 8,
"x1": 107.0,
"x2": 504.0,
"y1": 357.0,
"y2": 423.0
},
{
"page_number": 8,
"x1": 107.0,
"x2": 505.0,
"y1": 428.0,
"y2": 473.0
}
],
"keywords": [
"Transformer",
"Attention Mechanism",
"Machine Translation",
"Neural Networks",
"BLEU Score"
],
"keywords_citation": [
{
"page_number": 1,
"x1": 210.0,
"x2": 400.0,
"y1": 99.0,
"y2": 115.0
},
{
"page_number": 1,
"x1": 142.0,
"x2": 469.0,
"y1": 359.0,
"y2": 500.0
},
{
"page_number": 8,
"x1": 107.0,
"x2": 504.0,
"y1": 357.0,
"y2": 423.0
},
{
"page_number": 8,
"x1": 107.0,
"x2": 505.0,
"y1": 428.0,
"y2": 473.0
}
],
"sections": [
{
"heading": "1 Introduction",
"heading_citation": [],
"summary": "Recurrent neural networks and their variants have been the state of the art in sequence modeling. However, they have limitations in parallelization and computational efficiency. The Transformer model addresses these issues by using attention mechanisms instead of recurrence.",
"summary_citation": [
{
"page_number": 1,
"x1": 107.0,
"x2": 504.0,
"y1": 544.0,
"y2": 598.0
},
{
"page_number": 2,
"x1": 107.0,
"x2": 504.0,
"y1": 73.0,
"y2": 161.0
}
]
},
{
"heading": "2 Background",
"heading_citation": [],
"summary": "The Transformer model builds on previous work in reducing sequential computation, using attention mechanisms to model dependencies without regard to distance in sequences.",
"summary_citation": [
{
"page_number": 2,
"x1": 106.0,
"x2": 505.0,
"y1": 298.0,
"y2": 397.0
},
{
"page_number": 2,
"x1": 107.0,
"x2": 505.0,
"y1": 402.0,
"y2": 447.0
}
]
},
{
"heading": "3 Model Architecture",
"heading_citation": [],
"summary": "The Transformer uses a novel architecture based on self-attention and feed-forward networks, allowing for more parallelization and improved translation quality.",
"summary_citation": [
{
"page_number": 2,
"x1": 106.0,
"x2": 505.0,
"y1": 573.0,
"y2": 629.0
},
{
"page_number": 2,
"x1": 106.0,
"x2": 505.0,
"y1": 633.0,
"y2": 668.0
}
]
},
{
"heading": "4 Why Self-Attention",
"heading_citation": [],
"summary": "Self-attention layers offer advantages in computational complexity and learning long-range dependencies compared to recurrent and convolutional layers.",
"summary_citation": [
{
"page_number": 6,
"x1": 107.0,
"x2": 505.0,
"y1": 476.0,
"y2": 530.0
},
{
"page_number": 6,
"x1": 107.0,
"x2": 504.0,
"y1": 536.0,
"y2": 559.0
}
]
},
{
"heading": "5 Training",
"heading_citation": [],
"summary": "The Transformer model was trained on large datasets using efficient batching and hardware, achieving high performance with relatively low training costs.",
"summary_citation": [
{
"page_number": 7,
"x1": 106.0,
"x2": 250.0,
"y1": 304.0,
"y2": 316.0
},
{
"page_number": 7,
"x1": 106.0,
"x2": 505.0,
"y1": 324.0,
"y2": 401.0
}
]
},
{
"heading": "6 Results",
"heading_citation": [],
"summary": "The Transformer model outperforms previous state-of-the-art models in translation tasks, achieving higher BLEU scores with less training cost.",
"summary_citation": [
{
"page_number": 8,
"x1": 107.0,
"x2": 504.0,
"y1": 357.0,
"y2": 423.0
},
{
"page_number": 8,
"x1": 107.0,
"x2": 505.0,
"y1": 428.0,
"y2": 473.0
}
]
},
{
"heading": "7 Conclusion",
"heading_citation": [],
"summary": "The Transformer introduces a new paradigm in sequence transduction models, offering faster training and superior performance. Future work will explore its application to other modalities and tasks.",
"summary_citation": [
{
"page_number": 9,
"x1": 107.0,
"x2": 504.0,
"y1": 525.0,
"y2": 558.0
},
{
"page_number": 9,
"x1": 107.0,
"x2": 504.0,
"y1": 564.0,
"y2": 608.0
}
]
}
],
"title": "Attention Is All You Need",
"title_citation": []
}
},
"page_numbers": [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
11
],
"schema_name": "research_paper"
}
]
@@ -0,0 +1,27 @@
[project]
name = "context-engineering"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"crewai>=0.165.1",
"crewai-tools>=0.0.1",
"ipykernel>=6.30.1",
"openai>=1.99.9",
"pymilvus>=2.6.0",
"tensorlake>=0.2.44",
"voyageai>=0.3.4",
"zep-crewai>=1.0.0",
"python-dotenv>=1.0.0",
"requests>=2.31.0",
"pydantic>=2.0.0",
"firecrawl-py>=4.3.0",
"streamlit>=1.49.1",
"pyyaml>=6.0.2",
]
[dependency-groups]
dev = [
"ipykernel>=6.30.1",
]
@@ -0,0 +1,61 @@
import os
import yaml
from pathlib import Path
from typing import Dict, Any, Optional
class ConfigLoader:
"""Utility class for loading YAML configuration files"""
def __init__(self, config_root: Optional[str] = None):
if config_root is None:
project_root = Path(__file__).parent.parent.parent
self.config_root = project_root / "config"
else:
self.config_root = Path(config_root)
def load_agents_config(self, config_file: str = "research_agents.yaml") -> Dict[str, Any]:
"""Load agents configuration from YAML file"""
config_path = self.config_root / "agents" / config_file
return self._load_yaml_file(config_path)
def load_tasks_config(self, config_file: str = "research_tasks.yaml") -> Dict[str, Any]:
"""Load tasks configuration from YAML file"""
config_path = self.config_root / "tasks" / config_file
return self._load_yaml_file(config_path)
def _load_yaml_file(self, file_path: Path) -> Dict[str, Any]:
"""Load a YAML file and return its contents"""
try:
if not file_path.exists():
raise FileNotFoundError(f"Configuration file not found: {file_path}")
with open(file_path, 'r', encoding='utf-8') as file:
content = yaml.safe_load(file)
if content is None:
raise ValueError(f"Empty or invalid YAML file: {file_path}")
return content
except yaml.YAMLError as e:
raise ValueError(f"Error parsing YAML file {file_path}: {e}")
except Exception as e:
raise Exception(f"Error loading configuration file {file_path}: {e}")
def get_agent_config(self, agent_name: str, config_file: str = "research_agents.yaml") -> Dict[str, Any]:
"""Get configuration for a specific agent"""
agents_config = self.load_agents_config(config_file)
if agent_name not in agents_config:
raise KeyError(f"Agent '{agent_name}' not found in configuration file")
return agents_config[agent_name]
def get_task_config(self, task_name: str, config_file: str = "research_tasks.yaml") -> Dict[str, Any]:
"""Get configuration for a specific task"""
tasks_config = self.load_tasks_config(config_file)
if task_name not in tasks_config:
raise KeyError(f"Task '{task_name}' not found in configuration file")
return tasks_config[task_name]
@@ -0,0 +1,218 @@
import os
import json
from typing import Iterable, List, Dict, Any, Optional
from dotenv import load_dotenv
load_dotenv()
from tensorlake.documentai import (
DocumentAI,
ParsingOptions,
ChunkingStrategy,
TableOutputMode,
TableParsingFormat,
StructuredExtractionOptions
)
TENSORLAKE_API_KEY = os.getenv("TENSORLAKE_API_KEY")
RESEARCH_PAPER_SCHEMA = {
"type": "object",
"properties": {
"paper": {
"type": "object",
"properties": {
"title": {"type": "string"},
"authors": {
"type": "array",
"items": {"type": "string"}
},
"abstract": {"type": "string"},
"keywords": {
"type": "array",
"items": {"type": "string"}
},
"key_findings": {
"type": "array",
"items": {"type": "string"}
},
"sections": {
"type": "array",
"items": {
"type": "object",
"properties": {
"heading": {"type": "string"},
"summary": {"type": "string"}
},
"required": ["heading", "summary"]
}
}
},
"required": ["title", "authors", "abstract", "sections"]
}
},
"required": ["paper"]
}
class TensorLakeClient:
def __init__(self, api_key: Optional[str] = None):
self.doc_ai = DocumentAI(api_key=api_key or TENSORLAKE_API_KEY)
def list_uploaded_files(self):
try:
files_page = self.doc_ai.files()
print(f"TensorLake files found: {len(files_page.items)}")
for file_info in files_page.items:
print(f" - {file_info.name} (ID: {file_info.id}, Size: {file_info.file_size} bytes, Type: {file_info.mime_type})")
return files_page.items
except Exception as e:
print(f"Error listing TensorLake files: {e}")
return []
def verify_file_uploaded(self, file_id: str) -> bool:
try:
files = self.list_uploaded_files()
file_ids = [f.id for f in files]
exists = file_id in file_ids
print(f"File ID {file_id} {'exists' if exists else 'NOT FOUND'} in TensorLake")
return exists
except Exception as e:
print(f"Error verifying file {file_id}: {e}")
return False
def upload(self, paths: Iterable[str]) -> List[str]:
print("Files before upload:")
files_before = self.list_uploaded_files()
file_ids = []
for path in paths:
if not os.path.exists(path):
raise Exception(f"File does not exist: {path}")
file_size = os.path.getsize(path)
if file_size == 0:
raise Exception(f"File is empty: {path} (0 bytes)")
print(f"\nUploading file: {path} ({file_size} bytes)")
try:
fid = self.doc_ai.upload(path=path)
print(f"Upload successful, file_id: {fid}")
file_ids.append(fid)
except Exception as upload_error:
print(f"Upload failed for {path}: {upload_error}")
raise
print("\nFiles after upload:")
files_after = self.list_uploaded_files()
new_files = [f for f in files_after if f not in files_before]
if new_files:
print(f"{len(new_files)} new file(s) uploaded:")
for file_info in new_files:
print(f" - {file_info.name} (ID: {file_info.id})")
else:
print("No new files detected in TensorLake after upload")
return file_ids
def parse_structured(
self,
file_id: str,
json_schema: Dict[str, Any],
*,
page_range = None,
labels = None,
chunking_strategy = ChunkingStrategy.SECTION,
table_mode = TableOutputMode.MARKDOWN,
table_format = TableParsingFormat.TSR,
) -> str:
print(f"Using chunking strategy: {chunking_strategy}")
print(f"Using table mode: {table_mode}")
print(f"Schema name: research_paper")
structured_extraction_options = StructuredExtractionOptions(
schema_name="research_paper",
json_schema=json_schema,
provide_citations=True
)
parsing_options = ParsingOptions(
chunking_strategy=chunking_strategy,
table_output_mode=table_mode,
table_parsing_format=table_format,
)
if not self.verify_file_uploaded(file_id):
raise Exception(f"File ID {file_id} not found in TensorLake. Cannot proceed with parsing.")
print(f"Initiating parsing for file_id: {file_id}")
try:
parse_id = self.doc_ai.parse(
file_id,
page_range=page_range,
parsing_options=parsing_options,
structured_extraction_options=structured_extraction_options,
labels=labels or {}
)
print(f"Parsing initiated, parse_id: {parse_id}")
return parse_id
except Exception as parse_error:
print(f"Parsing initiation failed: {parse_error}")
raise
def get_result(self, parse_id: str) -> Dict[str, Any]:
print(f"Waiting for completion of parse_id: {parse_id}")
result = self.doc_ai.wait_for_completion(parse_id)
print(f"Parsing completed for parse_id: {parse_id}")
if result:
if hasattr(result, 'chunks'):
chunks = result.chunks
chunk_count = len(chunks) if chunks else 0
print(f"Number of chunks found: {chunk_count}")
if chunk_count > 0:
print(f"First chunk preview: {chunks[0].content[:100] if hasattr(chunks[0], 'content') else 'No content'}...")
else:
print("Result has no 'chunks' attribute")
else:
print("Result is None or empty")
return result
if __name__ == "__main__":
client = TensorLakeClient()
# Upload local documents
file_ids = client.upload([
"data/attention-is-all-you-need-Paper.pdf",
])
# Parse each with schema (and produce RAG chunks)
parse_ids = []
for fid in file_ids:
pid = client.parse_structured(
file_id=fid,
json_schema=RESEARCH_PAPER_SCHEMA,
page_range=None, # parse all pages
)
parse_ids.append(pid)
# Retrieve the parsed result (markdown chunks + schema JSON)
results = [client.get_result(pid) for pid in parse_ids]
# collect RAG chunks + structured paper metadata
rag_chunks, extracted_data = [], []
for res in results:
# markdown chunks for retrieval
for chunk in res.chunks:
rag_chunks.append({
"page": chunk.page_number,
"text": chunk.content,
})
# structured extraction schema
serializable_data = res.model_dump()
extracted_data.append(serializable_data)
@@ -0,0 +1,106 @@
import os
import json
from typing import Dict, Any, List, Optional
from openai import OpenAI
SYSTEM_PROMPT = """You are a research assistant that MUST ground answers in provided context.
Policy:
1) Use only the supplied CONTEXT and/or explicit SOURCE items.
2) If context is INSUFFICIENT to answer the QUESTION, return status=INSUFFICIENT_CONTEXT with what is missing.
3) If sufficient, answer concisely with citations (doc/page or URL) and a confidence score (01).
4) Never rely on parametric knowledge if it is not in the context.
5) Output MUST match the response schema exactly.
"""
RAG_TEMPLATE = (
"CONTEXT:\n{context}\n"
"---------------------\n"
"QUESTION:\n{query}\n\n"
"Task: Determine if the CONTEXT is sufficient to answer the QUESTION.\n"
"- If sufficient: produce a grounded answer with citations and confidence.\n"
"- If NOT sufficient: do NOT answer; return status=INSUFFICIENT_CONTEXT and list missing info.\n"
"Fill the structured fields only.\n"
)
# Structured Output schema
RESPONSE_SCHEMA = {
"type": "object",
"properties": {
"status": {"type": "string", "enum": ["OK", "INSUFFICIENT_CONTEXT"]},
"source_used": {"type": "string", "enum": ["MEMORY", "RAG", "WEB", "TOOL", "NONE"]},
"answer": {"type": "string"},
"citations": {
"type": "array",
"items": {
"type": "object",
"properties": {
"label": {"type": "string"},
"locator": {"type": "string"}
},
"required": ["label", "locator"],
"additionalProperties": False
}
},
"confidence": {"type": "number", "minimum": 0, "maximum": 1},
"missing": {"type": "array", "items": {"type": "string"}}
},
"required": ["status", "source_used", "answer", "citations", "confidence", "missing"],
"additionalProperties": False
}
class StructuredResponseGen:
def __init__(
self,
api_key: Optional[str] = None,
model: str = "gpt-4o-mini",
system_prompt: str = SYSTEM_PROMPT,
rag_template: str = RAG_TEMPLATE,
temperature: float = 0.2,
):
self.client = OpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY"))
self.model = model
self.system_prompt = system_prompt
self.rag_template = rag_template
self.temperature = temperature
def generate(
self,
*,
query: str,
context_blocks: List[str],
source_used: str = "RAG",
schema: Dict[str, Any] = RESPONSE_SCHEMA,
) -> Dict[str, Any]:
context = "\n\n".join(context_blocks).strip()
user_prompt = self.rag_template.format(context=context, query=query)
response = self.client.chat.completions.create(
model=self.model,
temperature=self.temperature,
messages=[
{"role": "system", "content": self.system_prompt},
{"role": "user", "content": user_prompt},
],
response_format={
"type": "json_schema",
"json_schema": {
"name": "research_briefing",
"schema": schema,
"strict": True # hard schema adherence
}
},
)
try:
output_text = response.choices[0].message.content
except Exception as e:
raise RuntimeError(f"Unexpected responses payload shape: {e}")
try:
data = json.loads(output_text)
except json.JSONDecodeError as e:
raise RuntimeError(f"Model did not return valid JSON: {e}\nRaw: {output_text[:400]}")
data["source_used"] = source_used
return data
@@ -0,0 +1,70 @@
import os
import time
from typing import Optional, Any, Dict
from zep_cloud.client import Zep
from zep_crewai import ZepUserStorage
from crewai.memory.external.external_memory import ExternalMemory
class ZepMemoryLayer:
def __init__(
self,
user_id: str,
thread_id: str,
mode: str = "summary",
indexing_wait_time: int = 10,
zep_api_key: Optional[str] = None,
):
self.zep_client = Zep(api_key=zep_api_key or os.getenv("ZEP_API_KEY"))
self.user_id = user_id
self.thread_id = thread_id
self.indexing_wait_time = indexing_wait_time
try:
self.zep_client.user.get(self.user_id)
except:
self.zep_client.user.add(user_id=self.user_id)
# Create new session by first deleting the previous one
self.zep_client.thread.delete(self.thread_id)
self.zep_client.thread.create(thread_id=self.thread_id, user_id=self.user_id)
self.user_storage = ZepUserStorage(
client=self.zep_client,
user_id=self.user_id,
thread_id=self.thread_id,
mode=mode,
)
self.external_memory = ExternalMemory(storage=self.user_storage)
def as_external_memory(self) -> ExternalMemory:
return self.external_memory
def save_user_message(self, text: str, name: Optional[str] = None, **meta: Any) -> None:
self.external_memory.save(
text,
metadata={"type": "message", "role": "user", "name": name or "User", **meta},
)
def save_assistant_message(self, text: str, name: Optional[str] = None, **meta: Any) -> None:
self.external_memory.save(
text,
metadata={"type": "message", "role": "assistant", "name": name or "Assistant", **meta},
)
def save_preferences(self, prefs: Dict[str, Any]) -> None:
self.external_memory.save(
str({"preferences": prefs}),
metadata={"type": "json", "category": "preferences"},
)
def wait_for_indexing(self) -> None:
time.sleep(self.indexing_wait_time)
def get_context_block(self) -> str:
"""
Fetch the full user context block from Zep memory.
Returns a string of concatenated memory facts (Zep's internal summary).
"""
memory = self.zep_client.thread.get_user_context(thread_id=self.thread_id)
return memory.context if memory and memory.context else ""
@@ -0,0 +1,48 @@
import os
import voyageai
from typing import List, Optional, Literal
from dotenv import load_dotenv
load_dotenv()
VOYAGE_API_KEY = os.getenv("VOYAGE_API_KEY")
class ContextualizedEmbeddings:
def __init__(self, api_key: Optional[str] = None, model: str = "voyage-context-3"):
self.client = voyageai.Client(api_key=api_key or VOYAGE_API_KEY)
self.model = model
def embed_document_chunks(
self,
docs_chunks: List[List[str]],
*,
output_dimension = 1024,
output_dtype = "float",
) -> List[List[List[float]]]:
resp = self.client.contextualized_embed(
inputs=docs_chunks,
model=self.model,
input_type="document",
output_dimension=output_dimension,
output_dtype=output_dtype,
)
return [r.embeddings for r in resp.results]
def embed_query(
self,
query,
*,
output_dimension = None,
output_dtype: Literal["float", "int8", "uint8", "binary", "ubinary"] = "float",
) -> List[float]:
resp = self.client.contextualized_embed(
inputs=[[query]],
model=self.model,
input_type="query",
output_dimension=output_dimension,
output_dtype=output_dtype,
)
return resp.results[0].embeddings[0]
@@ -0,0 +1,137 @@
import os
from typing import List, Dict, Any, Optional
from src.document_processing import TensorLakeClient, RESEARCH_PAPER_SCHEMA
from src.rag.embeddings import ContextualizedEmbeddings
from src.rag.retriever import MilvusVectorDB
from src.generation import StructuredResponseGen
class RAGPipeline:
"""Unified RAG pipeline combining document parsing, embeddings, and retrieval"""
def __init__(
self,
tensorlake_api_key: Optional[str] = None,
voyage_api_key: Optional[str] = None,
openai_api_key: Optional[str] = None,
milvus_db_path: str = "milvus_lite.db",
collection_name: str = "research_assistant"
):
self.doc_parser = TensorLakeClient(api_key=tensorlake_api_key)
self.embeddings = ContextualizedEmbeddings(api_key=voyage_api_key)
self.vector_db = MilvusVectorDB(db_path=milvus_db_path, collection_name=collection_name)
self.generator = StructuredResponseGen(api_key=openai_api_key)
def process_documents(self, document_paths: List[str]) -> Dict[str, Any]:
results = {
"processed_docs": [],
"total_chunks": 0,
"structured_data": []
}
# Upload and parse documents
file_ids = self.doc_parser.upload(document_paths)
for i, (path, file_id) in enumerate(zip(document_paths, file_ids)):
parse_id = self.doc_parser.parse_structured(
file_id=file_id,
json_schema=RESEARCH_PAPER_SCHEMA,
labels={"source": path, "doc_index": i}
)
parse_result = self.doc_parser.get_result(parse_id)
if parse_result is None:
raise Exception(f"TensorLake parsing failed for {path}: No result returned")
# Extract chunks and structured data
chunks = []
if hasattr(parse_result, 'chunks') and parse_result.chunks:
for chunk in parse_result.chunks:
if chunk and hasattr(chunk, 'content') and hasattr(chunk, 'page_number'):
chunks.append({
"page": chunk.page_number,
"text": chunk.content,
"source": path
})
else:
raise Exception(f"TensorLake parsing failed for {path}: No chunks found in result")
if not chunks:
raise Exception(f"TensorLake parsing failed for {path}: No valid chunks extracted")
# Generate contextualized embeddings
chunk_texts = [[chunk["text"] for chunk in chunks]]
embeddings_result = self.embeddings.embed_document_chunks(chunk_texts)
if not embeddings_result or len(embeddings_result) == 0:
raise Exception(f"Embedding generation failed for {path}: No embeddings returned")
chunk_embeddings = embeddings_result[0]
if not chunk_embeddings or len(chunk_embeddings) != len(chunks):
raise Exception(f"Embedding generation failed for {path}: Chunk count mismatch - {len(chunks)} chunks but {len(chunk_embeddings) if chunk_embeddings else 0} embeddings")
# Prepare metadata for each chunk
chunk_metadata = []
for i, chunk in enumerate(chunks):
metadata = {
"page_number": chunk.get("page", 0),
"chunk_index": i,
"source_file": chunk.get("source", path)
}
chunk_metadata.append(metadata)
# Store in vector database with metadata
self.vector_db.insert(
chunks=[chunk["text"] for chunk in chunks],
embeddings=chunk_embeddings,
metadata=chunk_metadata
)
results["processed_docs"].append({
"path": path,
"file_id": file_id,
"chunks_count": len(chunks),
"structured_data": parse_result.model_dump()
})
results["total_chunks"] += len(chunks)
results["structured_data"].append(parse_result.model_dump())
return results
def retrieve_context(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
query_embedding = self.embeddings.embed_query(query)
# Search vector database
search_results = self.vector_db.search(
query_embedding=query_embedding,
limit=top_k
)
return search_results
def generate_response(
self,
query: str,
context: List[Dict[str, Any]],
source_used: str = "RAG"
) -> Dict[str, Any]:
context_blocks = [result["text"] for result in context]
response = self.generator.generate(
query=query,
context_blocks=context_blocks,
source_used=source_used
)
return response
def query(self, query: str, top_k: int = 3) -> Dict[str, Any]:
context_results = self.retrieve_context(query, top_k=top_k)
response = self.generate_response(query, context_results)
# Add retrieval metadata for citations
response["retrieval_metadata"] = {
"retrieved_chunks": len(context_results),
"top_scores": [r["score"] for r in context_results]
}
return response
@@ -0,0 +1,101 @@
from typing import List, Dict, Any
from pymilvus import MilvusClient, DataType
class MilvusVectorDB:
def __init__(self, db_path: str = "milvus_lite.db", collection_name: str = "research_assistant"):
self.client = MilvusClient(db_path)
self.collection_name = collection_name
self._ensure_collection()
def _ensure_collection(self, dim: int = 1024):
if self.client.has_collection(collection_name=self.collection_name):
print(f"Dropping existing collection: {self.collection_name}")
self.client.drop_collection(collection_name=self.collection_name)
schema = self.client.create_schema(
auto_id=True,
enable_dynamic_fields=True,
)
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=dim)
schema.add_field("text", DataType.VARCHAR, max_length=65535)
schema.add_field("page_number", DataType.INT64)
schema.add_field("chunk_index", DataType.INT64)
schema.add_field("source_file", DataType.VARCHAR, max_length=500)
index_params = self.client.prepare_index_params()
index_params.add_index("embedding", index_type="IVF_FLAT", metric_type="COSINE")
self.client.create_collection(
collection_name=self.collection_name,
schema=schema,
index_params=index_params
)
def insert(self, chunks: List[str], embeddings: List[List[float]], metadata: List[Dict[str, Any]] = None):
assert len(chunks) == len(embeddings), "Mismatch between chunks and embeddings"
if metadata:
assert len(chunks) == len(metadata), "Mismatch between chunks and metadata"
data = []
for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
entry = {
"text": chunk,
"embedding": emb,
}
if metadata and i < len(metadata):
meta = metadata[i]
entry["page_number"] = meta.get("page_number", 0)
entry["chunk_index"] = meta.get("chunk_index", i)
entry["source_file"] = meta.get("source_file", "unknown")
else:
entry["page_number"] = 0
entry["chunk_index"] = i
entry["source_file"] = "unknown"
data.append(entry)
self.client.insert(
collection_name=self.collection_name,
data=data
)
self.client.flush(collection_name=self.collection_name)
def get_collection_count(self) -> int:
try:
stats = self.client.get_collection_stats(collection_name=self.collection_name)
return stats.get('row_count', 0)
except:
return 0
def search(
self,
query_embedding: List[float],
limit: int = 3,
nprobe: int = 10,
metric: str = "COSINE"
) -> List[Dict[str, Any]]:
search_params = {"metric_type": metric, "params": {"nprobe": nprobe}}
results = self.client.search(
collection_name=self.collection_name,
data=[query_embedding],
anns_field="embedding",
search_params=search_params,
limit=limit,
output_fields=["text", "page_number", "chunk_index", "source_file"],
)
hits = []
for hit in results[0]:
hits.append({
"text": hit.entity.get("text"),
"score": hit.score,
"page_number": hit.entity.get("page_number", 0),
"chunk_index": hit.entity.get("chunk_index", 0),
"source_file": hit.entity.get("source_file", "unknown")
})
return hits
@@ -0,0 +1,177 @@
import json
import requests
import xml.etree.ElementTree as ET
from typing import Dict, Any, List, Optional, Type
from pydantic import BaseModel, Field
from crewai.tools import BaseTool
class ArxivAPIInput(BaseModel):
"""Input schema for ArXiv API tool"""
query: str = Field(..., description="The research query to search for papers.")
search_field: str = Field(default="all", description="Field to search in: 'all', 'title', 'author', 'abstract', 'category'")
category: Optional[str] = Field(default=None, description="ArXiv category filter (e.g., 'cs.AI', 'stat.ML', 'physics')")
author: Optional[str] = Field(default=None, description="Author name to filter by")
max_results: int = Field(default=5, description="Maximum number of papers to return (1-50)")
class ArxivTool(BaseTool):
name: str = "arxiv_search"
description: str = "Search ArXiv for academic papers related to your research query. Can filter by category, author, and search specific fields."
args_schema: Type[BaseModel] = ArxivAPIInput
def _run(self, query: str, search_field: str = "all", category: Optional[str] = None,
author: Optional[str] = None, max_results: int = 5) -> str:
try:
# Build ArXiv query
search_query = self._build_arxiv_query(query, search_field, category, author)
# ArXiv API endpoint
base_url = "http://export.arxiv.org/api/query"
params = {
"search_query": search_query,
"start": 0,
"max_results": max_results,
"sortBy": "relevance",
"sortOrder": "descending"
}
# Make API request
response = requests.get(base_url, params=params, timeout=30)
response.raise_for_status()
# Parse XML response
papers = self._parse_arxiv_response(response.text)
if not papers:
return json.dumps({
"status": "INSUFFICIENT_CONTEXT",
"source_used": "ARXIV",
"answer": f"No papers found for query: '{query}'",
"citations": [],
"confidence": 0.0,
"search_parameters": {
"query": query,
"search_field": search_field,
"category": category,
"author": author,
"max_results": max_results
}
})
answer_parts = []
citations = []
for i, paper in enumerate(papers):
answer_parts.append(
f"**{i+1}. {paper['title']}**\n"
f"Authors: {paper['authors']}\n"
f"Category: {paper['category']}\n"
f"Published: {paper['published']}\n"
f"Abstract: {paper['abstract'][:300]}...\n"
f"URL: {paper['url']}"
)
citations.append({
"label": f"{paper['title']} ({paper['published']})",
"locator": paper['url']
})
answer = f"Found {len(papers)} relevant papers:\n\n" + "\n\n---\n\n".join(answer_parts)
return json.dumps({
"status": "OK",
"source_used": "ARXIV",
"answer": answer,
"citations": citations,
"confidence": 0.92,
"search_parameters": {
"query": query,
"search_field": search_field,
"category": category,
"author": author,
"max_results": max_results
},
"papers_found": len(papers)
})
except Exception as e:
return json.dumps({
"status": "ERROR",
"source_used": "ARXIV",
"answer": f"ArXiv search failed: {str(e)}",
"citations": [],
"confidence": 0.0,
"error": str(e)
})
def _build_arxiv_query(self, query: str, search_field: str, category: Optional[str], author: Optional[str]) -> str:
parts = []
if search_field == "title":
parts.append(f'ti:"{query}"')
elif search_field == "author":
parts.append(f'au:"{query}"')
elif search_field == "abstract":
parts.append(f'abs:"{query}"')
elif search_field == "category":
parts.append(f'cat:"{query}"')
else: # search_field == "all"
parts.append(f'all:"{query}"')
if category:
parts.append(f'cat:{category}')
if author:
parts.append(f'au:"{author}"')
return " AND ".join(parts)
def _parse_arxiv_response(self, xml_content: str) -> List[Dict[str, Any]]:
try:
root = ET.fromstring(xml_content)
namespaces = {
'atom': 'http://www.w3.org/2005/Atom',
'arxiv': 'http://arxiv.org/schemas/atom'
}
papers = []
entries = root.findall('atom:entry', namespaces)
for entry in entries:
title = entry.find('atom:title', namespaces)
title_text = title.text.strip().replace('\n', ' ') if title is not None else "No title"
authors = []
author_elements = entry.findall('atom:author', namespaces)
for author in author_elements:
name = author.find('atom:name', namespaces)
if name is not None:
authors.append(name.text)
authors_text = ", ".join(authors) if authors else "Unknown authors"
summary = entry.find('atom:summary', namespaces)
abstract = summary.text.strip().replace('\n', ' ') if summary is not None else "No abstract"
link = entry.find('atom:id', namespaces)
url = link.text if link is not None else ""
published = entry.find('atom:published', namespaces)
pub_date = published.text[:10] if published is not None else "Unknown date"
category_elem = entry.find('arxiv:primary_category', namespaces)
if category_elem is None:
category_elem = entry.find('atom:category', namespaces)
category = category_elem.get('term') if category_elem is not None else "Unknown category"
papers.append({
"title": title_text,
"authors": authors_text,
"abstract": abstract,
"url": url,
"published": pub_date,
"category": category
})
return papers
except ET.ParseError as e:
raise Exception(f"Failed to parse ArXiv XML response: {e}")
except Exception as e:
raise Exception(f"Error processing ArXiv response: {e}")
@@ -0,0 +1,51 @@
import json
from typing import Type
from pydantic import BaseModel, Field
from crewai.tools import BaseTool
from src.memory import ZepMemoryLayer
class MemoryInput(BaseModel):
"""Input schema for memory search tool"""
query: str = Field(..., description="The search query for memory retrieval.")
class MemoryTool(BaseTool):
name: str = "memory_search"
description: str = "Retrieve relevant information from conversation history and user preferences"
memory_layer: ZepMemoryLayer = Field(..., description="Zep memory layer instance")
args_schema: Type[BaseModel] = MemoryInput
def _run(self, query: str) -> str:
try:
context = self.memory_layer.get_context_block()
if context:
result = {
"status": "OK",
"source_used": "MEMORY",
"answer": f"Retrieved relevant context from previous conversations: {context}",
"citations": [{"label": "Conversation History", "locator": "zep:memory"}],
"confidence": 0.98,
"context": context
}
else:
result = {
"status": "INSUFFICIENT_CONTEXT",
"source_used": "MEMORY",
"answer": "No relevant conversation history found",
"citations": [],
"confidence": 0.0,
"context": ""
}
return json.dumps(result, indent=2)
except Exception as e:
return json.dumps({
"status": "ERROR",
"source_used": "MEMORY",
"answer": f"Memory retrieval failed: {str(e)}",
"citations": [],
"confidence": 0.0,
"context": ""
})
@@ -0,0 +1,166 @@
import os
import json
from typing import Dict, Any, List, Optional, Type
from pydantic import BaseModel, Field
from crewai.tools import BaseTool
class RAGInput(BaseModel):
"""Input schema for RAG search tool"""
query: str = Field(..., description="The search query for retrieval.")
top_k: int = Field(default=3, description="Maximum number of retrieved results to fetch.")
document_paths: List[str] = Field(default=None, description="Optional list of document paths to load. Only needed if no documents are already loaded in the vector database.")
class RAGTool(BaseTool):
name: str = "rag_search"
description: str = "Search through research documents for relevant information."
rag_pipeline: Any = Field(..., description="RAG pipeline instance")
args_schema: Type[BaseModel] = RAGInput
def _run(self, query: str, top_k: int = 3, document_paths: List[str] = None):
try:
doc_count = self.rag_pipeline.vector_db.get_collection_count()
if doc_count == 0:
if not document_paths:
return json.dumps({
"status": "INSUFFICIENT_CONTEXT",
"source_used": "RAG",
"answer": "No documents have been loaded into the RAG system. Please provide document_paths to load documents first, or ensure documents have been previously loaded.",
"citations": [],
"confidence": 0.0,
"retrieval_metadata": {
"retrieved_chunks": 0,
"top_scores": [],
"document_count": 0
}
})
# load documents
load_result = self._load_documents(document_paths)
if load_result["status"] == "ERROR":
return json.dumps(load_result, indent=2)
doc_count = self.rag_pipeline.vector_db.get_collection_count()
if doc_count == 0:
return json.dumps({
"status": "INSUFFICIENT_CONTEXT",
"source_used": "RAG",
"answer": "Failed to load documents into the RAG system.",
"citations": [],
"confidence": 0.0,
"retrieval_metadata": {
"retrieved_chunks": 0,
"top_scores": [],
"document_count": 0
}
})
# Retrieve relevant context (no generation)
context_results = self.rag_pipeline.retrieve_context(query, top_k=top_k)
if not context_results:
return json.dumps({
"status": "INSUFFICIENT_CONTEXT",
"source_used": "RAG",
"answer": f"No relevant context found for query: '{query}'",
"citations": [],
"confidence": 0.0,
"retrieval_metadata": {
"retrieved_chunks": 0,
"top_scores": [],
"document_count": doc_count
}
})
context_blocks = []
citations = []
for i, result in enumerate(context_results):
chunk_text = result.get("text", "")
score = result.get("score", 0.0)
page_number = result.get("page_number", 0)
chunk_index = result.get("chunk_index", i)
source_file = result.get("source_file", "unknown")
filename = source_file.split("/")[-1] if source_file != "unknown" else "unknown"
context_blocks.append(f"**Context {i+1} (Score: {score:.3f}, Page {page_number}, Chunk {chunk_index})**\n{chunk_text[:500]}...")
citations.append({
"label": f"{filename} - Page {page_number}, Chunk {chunk_index}",
"locator": f"page_{page_number}_chunk_{chunk_index}",
"page_number": page_number,
"chunk_index": chunk_index,
"source_file": filename,
"score": score,
"content": chunk_text
})
answer = "\n\n".join(context_blocks)
return json.dumps({
"status": "OK",
"source_used": "RAG",
"answer": f"Retrieved {len(context_results)} relevant context chunks:\n\n{answer}",
"citations": citations,
"confidence": max([r.get("score", 0.0) for r in context_results]),
"retrieval_metadata": {
"retrieved_chunks": len(context_results),
"top_scores": [r.get("score", 0.0) for r in context_results],
"document_count": doc_count
},
"raw_context": context_results # Include raw context for further processing
})
except Exception as e:
return json.dumps({
"status": "ERROR",
"source_used": "RAG",
"answer": f"RAG search failed: {str(e)}",
"citations": [],
"confidence": 0.0,
"retrieval_metadata": {
"retrieved_chunks": 0,
"top_scores": [],
"document_count": 0
},
"error": str(e)
})
def _load_documents(self, document_paths: List[str]) -> Dict[str, Any]:
"""Helper method to load documents into the RAG pipeline"""
try:
if not document_paths:
return {
"status": "ERROR",
"source_used": "RAG",
"answer": "No document paths provided",
"citations": [],
"confidence": 0.0
}
missing_files = [path for path in document_paths if not os.path.exists(path)]
if missing_files:
return {
"status": "ERROR",
"source_used": "RAG",
"answer": f"Missing files: {missing_files}",
"citations": [],
"confidence": 0.0
}
results = self.rag_pipeline.process_documents(document_paths)
return {
"status": "OK",
"source_used": "RAG",
"answer": f"Successfully processed {len(results['processed_docs'])} documents with {results['total_chunks']} total chunks.",
"citations": [{"label": f"Processed: {doc['path']}", "locator": doc['path']} for doc in results['processed_docs']],
"confidence": 0.96
}
except Exception as e:
return {
"status": "ERROR",
"source_used": "RAG",
"answer": f"Document processing failed: {str(e)}",
"citations": [],
"confidence": 0.0,
"error": str(e)
}
@@ -0,0 +1,111 @@
import os
import json
from typing import Type
from pydantic import BaseModel, Field
from crewai.tools import BaseTool
from firecrawl import Firecrawl
class WebSearchInput(BaseModel):
"""Input schema for web search tool"""
query: str = Field(..., description="The search query for web search.")
limit: int = Field(default=3, description="Maximum number of search results to return.")
class FirecrawlSearchTool(BaseTool):
name: str = "firecrawl_web_search"
description: str = "Search the web for recent information and developments on research topics"
api_key: str = Field(..., description="Firecrawl API key")
args_schema: Type[BaseModel] = WebSearchInput
def _run(self, query: str, limit: int = 3) -> str:
try:
if not self.api_key:
return json.dumps({
"status": "ERROR",
"source_used": "WEB",
"answer": "Web search unavailable - API key not configured.",
"citations": [],
"confidence": 0.0,
"error": "Missing API key"
})
# Initialize Firecrawl app and perform search
app = Firecrawl(api_key=self.api_key)
response = app.search(query, limit=limit)
results_list = getattr(response, "web", None)
if not isinstance(results_list, list) or not results_list:
return json.dumps({
"status": "INSUFFICIENT_CONTEXT",
"source_used": "WEB",
"answer": "No relevant web search results found.",
"citations": [],
"confidence": 0.0,
"search_results": []
})
search_results = []
citations = []
for result in results_list:
try:
title = getattr(result, "title", "No title") or "No title"
url = getattr(result, "url", "") or ""
content = getattr(result, "description", "") or ""
category = getattr(result, "category", None)
# Truncate content for readability
snippet = content[:1000] + "..." if len(content) > 1000 else content
if not snippet:
snippet = "[no description available]"
search_results.append({
"title": title,
"url": url,
"content": snippet,
"category": category
})
citations.append({
"label": title,
"locator": url
})
except Exception as e:
continue
if search_results:
answer_parts = []
for result in search_results:
answer_parts.append(
f"**{result['title']}**\n"
f"URL: {result['url']}\n"
f"Content: {result['content'][:500]}..."
)
answer = "\n\n---\n\n".join(answer_parts)
return json.dumps({
"status": "OK",
"source_used": "WEB",
"answer": answer,
"citations": citations,
"confidence": 0.97,
"search_results": search_results
})
else:
return json.dumps({
"status": "INSUFFICIENT_CONTEXT",
"source_used": "WEB",
"answer": "No relevant web search results found.",
"citations": [],
"confidence": 0.0,
"search_results": []
})
except Exception as e:
return json.dumps({
"status": "ERROR",
"source_used": "WEB",
"answer": f"Web search unavailable due to technical issues: {str(e)}",
"citations": [],
"confidence": 0.0,
"error": str(e)
})
@@ -0,0 +1,81 @@
import os
from crewai import Agent
from typing import Optional
from src.config import ConfigLoader
from src.tools import (
RAGTool,
MemoryTool,
ArxivTool,
FirecrawlSearchTool
)
from src.rag import RAGPipeline
from src.memory import ZepMemoryLayer
class Agents:
"""Class for creating agents from configuration files"""
def __init__(self, config_loader: Optional[ConfigLoader] = None):
self.config_loader = config_loader or ConfigLoader()
def create_rag_agent(self, rag_pipeline: RAGPipeline) -> Agent:
config = self.config_loader.get_agent_config("rag_agent")
rag_tool = RAGTool(rag_pipeline=rag_pipeline)
return Agent(
role=config["role"],
goal=config["goal"],
backstory=config["backstory"],
tools=[rag_tool],
verbose=config.get("verbose", True)
)
def create_memory_agent(self, memory_layer: ZepMemoryLayer) -> Agent:
config = self.config_loader.get_agent_config("memory_agent")
memory_tool = MemoryTool(memory_layer=memory_layer)
return Agent(
role=config["role"],
goal=config["goal"],
backstory=config["backstory"],
tools=[memory_tool],
verbose=config.get("verbose", True)
)
def create_web_search_agent(self, firecrawl_api_key: str) -> Agent:
config = self.config_loader.get_agent_config("web_search_agent")
web_search_tool = FirecrawlSearchTool(api_key=firecrawl_api_key)
return Agent(
role=config["role"],
goal=config["goal"],
backstory=config["backstory"],
tools=[web_search_tool],
verbose=config.get("verbose", True)
)
def create_arxiv_agent(self) -> Agent:
config = self.config_loader.get_agent_config("arxiv_agent")
arxiv_tool = ArxivTool()
return Agent(
role=config["role"],
goal=config["goal"],
backstory=config["backstory"],
tools=[arxiv_tool],
verbose=config.get("verbose", True)
)
def create_evaluator_agent(self) -> Agent:
config = self.config_loader.get_agent_config("evaluator_agent")
return Agent(
role=config["role"],
goal=config["goal"],
backstory=config["backstory"],
verbose=config.get("verbose", True)
)
def create_synthesizer_agent(self) -> Agent:
config = self.config_loader.get_agent_config("synthesizer_agent")
return Agent(
role=config["role"],
goal=config["goal"],
backstory=config["backstory"],
verbose=config.get("verbose", True)
)
@@ -0,0 +1,246 @@
import os
import json
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field
from crewai import Crew, Task
from crewai.flow.flow import Flow, listen, start
from src.rag import RAGPipeline
from src.memory import ZepMemoryLayer
from .agents import Agents
from .tasks import Tasks
class ResearchAssistantState(BaseModel):
query: str = ""
user_id: str = "default_user"
thread_id: str = "default_thread"
class ContextEvaluationResult(BaseModel):
"""Pydantic schema for context evaluation agent output validation"""
relevant_sources: List[str] = Field(
...,
description="List of source names that are relevant to the query (e.g., 'RAG', 'Memory', 'Web', 'ArXiv')"
)
filtered_context: Dict[str, Any] = Field(
...,
description="Dictionary containing only relevant information from each source, keyed by source name"
)
relevance_scores: Dict[str, float] = Field(
...,
description="Confidence scores (0-1) for each source's relevance to the query"
)
reasoning: str = Field(
...,
description="Brief explanation of filtering decisions and why certain sources were included/excluded"
)
class Config:
schema_extra = {
"example": {
"relevant_sources": ["RAG", "Web"],
"filtered_context": {
"RAG": {
"status": "OK",
"answer": "Relevant document context...",
"citations": [{"label": "Paper Title", "locator": "chunk_1"}]
},
"Web": {
"status": "OK",
"answer": "Recent web information...",
"citations": [{"label": "Article Title", "locator": "https://example.com"}]
}
},
"relevance_scores": {
"RAG": 0.95,
"Memory": 0.3,
"Web": 0.85,
"ArXiv": 0.1
},
"reasoning": "RAG and Web sources contain highly relevant information for the query. Memory has some context but low relevance. ArXiv results don't match the specific query focus."
}
}
class ResearchAssistantFlow(Flow[ResearchAssistantState]):
def __init__(
self,
tensorlake_api_key: Optional[str] = None,
voyage_api_key: Optional[str] = None,
openai_api_key: Optional[str] = None,
zep_api_key: Optional[str] = None,
firecrawl_api_key: Optional[str] = None,
milvus_db_path: str = "milvus_lite.db",
):
super().__init__()
self.rag_pipeline = RAGPipeline(
tensorlake_api_key=tensorlake_api_key,
voyage_api_key=voyage_api_key,
openai_api_key=openai_api_key,
milvus_db_path=milvus_db_path
)
self.memory_layer = ZepMemoryLayer(
user_id=self.state.user_id,
thread_id=self.state.thread_id,
zep_api_key=zep_api_key
)
# Initialize tasks and agents
self.tasks = Tasks()
self.agents = Agents()
# Create agents
self.rag_agent = self.agents.create_rag_agent(self.rag_pipeline)
self.memory_agent = self.agents.create_memory_agent(self.memory_layer)
self.web_search_agent = self.agents.create_web_search_agent(firecrawl_api_key or os.getenv("FIRECRAWL_API_KEY"))
self.tool_calling_agent = self.agents.create_arxiv_agent()
self.evaluator_agent = self.agents.create_evaluator_agent()
self.synthesizer_agent = self.agents.create_synthesizer_agent()
@start()
def process_query(self) -> Dict[str, Any]:
query = self.state.query
# Save user query to memory
summarized_query = self._summarize_for_memory(query, max_length=1500)
self.memory_layer.save_user_message(summarized_query)
return {
"query": query,
"status": "processing"
}
@listen(process_query)
def gather_context_from_all_sources(self, flow_state: Dict[str, Any]) -> Dict[str, Any]:
query = flow_state["query"]
# Create tasks for each agent
rag_task = self.tasks.create_rag_search_task(query, self.rag_agent)
memory_task = self.tasks.create_memory_retrieval_task(query, self.memory_agent)
web_search_task = self.tasks.create_web_search_task(query, self.web_search_agent)
tool_calling_task = self.tasks.create_arxiv_search_task(query, self.tool_calling_agent)
context_crew = Crew(
agents=[self.rag_agent, self.memory_agent, self.web_search_agent, self.tool_calling_agent],
tasks=[rag_task, memory_task, web_search_task, tool_calling_task],
verbose=True
)
results = context_crew.kickoff()
# Parse results from each agent
context_sources = {
"rag_result": self._parse_agent_result(results.tasks_output[0].raw),
"memory_result": self._parse_agent_result(results.tasks_output[1].raw),
"web_result": self._parse_agent_result(results.tasks_output[2].raw),
"tool_result": self._parse_agent_result(results.tasks_output[3].raw)
}
return {
**flow_state,
"context_sources": context_sources,
"raw_results": [task.raw for task in results.tasks_output]
}
@listen(gather_context_from_all_sources)
def evaluate_context_relevance(self, flow_state: Dict[str, Any]) -> Dict[str, Any]:
query = flow_state["query"]
context_sources = flow_state["context_sources"]
evaluation_task = self.tasks.create_context_evaluation_task(
query, context_sources, self.evaluator_agent, ContextEvaluationResult
)
evaluation_crew = Crew(
agents=[self.evaluator_agent],
tasks=[evaluation_task],
verbose=True
)
evaluation_result = evaluation_crew.kickoff()
evaluation_output = evaluation_result.tasks_output[0].pydantic
if isinstance(evaluation_output, ContextEvaluationResult):
filtered_context = evaluation_output.filtered_context
evaluation_data = evaluation_output.model_dump()
else:
print("Pydantic output not available, falling back to raw parsing")
filtered_context = self._parse_agent_result(evaluation_result.tasks_output[0].raw)
evaluation_data = {"raw_fallback": evaluation_result.tasks_output[0].raw}
return {
**flow_state,
"filtered_context": filtered_context,
"evaluation_result": evaluation_data,
"evaluation_raw": evaluation_result.tasks_output[0].raw
}
@listen(evaluate_context_relevance)
def synthesize_final_response(self, flow_state: Dict[str, Any]) -> Dict[str, Any]:
query = flow_state["query"]
filtered_context = flow_state["filtered_context"]
synthesis_task = self.tasks.create_synthesis_task(
query, filtered_context, self.synthesizer_agent
)
synthesis_crew = Crew(
agents=[self.synthesizer_agent],
tasks=[synthesis_task],
verbose=True
)
synthesis_result = synthesis_crew.kickoff()
final_response = synthesis_result.tasks_output[0].raw
# Save summarized assistant response to memory
summarized_response = self._summarize_for_memory(final_response)
self.memory_layer.save_assistant_message(summarized_response)
return {
**flow_state,
"final_response": final_response,
"synthesis_raw": final_response,
"status": "completed"
}
def _parse_agent_result(self, raw_result: str) -> Dict[str, Any]:
try:
return json.loads(raw_result)
except json.JSONDecodeError:
return {
"status": "OK",
"source_used": "UNKNOWN",
"answer": raw_result,
"citations": [],
"confidence": 0.5
}
def _summarize_for_memory(self, response: str, max_length: int = 2000) -> str:
if len(response) <= max_length:
return response
truncated = response[:max_length]
last_period = truncated.rfind('.')
last_exclamation = truncated.rfind('!')
last_question = truncated.rfind('?')
last_sentence_end = max(last_period, last_exclamation, last_question)
if last_sentence_end > max_length * 0.7:
return truncated[:last_sentence_end + 1] + " [Response truncated for memory storage]"
else:
last_space = truncated.rfind(' ')
if last_space > max_length * 0.8:
return truncated[:last_space] + "... [Response truncated for memory storage]"
else:
return truncated + "... [Response truncated for memory storage]"
def process_documents(self, document_paths: List[str]) -> Dict[str, Any]:
return self.rag_pipeline.process_documents(document_paths)
def create_research_assistant_flow(**kwargs) -> ResearchAssistantFlow:
return ResearchAssistantFlow(**kwargs)
@@ -0,0 +1,74 @@
import json
from crewai import Task
from typing import Dict, Any, Optional
from src.config import ConfigLoader
class Tasks:
def __init__(self, config_loader: Optional[ConfigLoader] = None):
self.config_loader = config_loader or ConfigLoader()
def create_rag_search_task(self, query: str, agent) -> Task:
config = self.config_loader.get_task_config("rag_search_task")
return Task(
description=config["description"].format(query=query),
expected_output=config["expected_output"],
agent=agent
)
def create_memory_retrieval_task(self, query: str, agent) -> Task:
config = self.config_loader.get_task_config("memory_retrieval_task")
return Task(
description=config["description"].format(query=query),
expected_output=config["expected_output"],
agent=agent
)
def create_web_search_task(self, query: str, agent) -> Task:
config = self.config_loader.get_task_config("web_search_task")
return Task(
description=config["description"].format(query=query),
expected_output=config["expected_output"],
agent=agent
)
def create_arxiv_search_task(self, query: str, agent) -> Task:
config = self.config_loader.get_task_config("arxiv_search_task")
return Task(
description=config["description"].format(query=query),
expected_output=config["expected_output"],
agent=agent
)
def create_context_evaluation_task(self, query: str, context_sources: Dict[str, Any], agent, output_pydantic=None) -> Task:
config = self.config_loader.get_task_config("context_evaluation_task")
formatted_description = config["description"].format(
query=query,
rag_result=json.dumps(context_sources.get('rag_result', {}), indent=2),
memory_result=json.dumps(context_sources.get('memory_result', {}), indent=2),
web_result=json.dumps(context_sources.get('web_result', {}), indent=2),
tool_result=json.dumps(context_sources.get('tool_result', {}), indent=2)
)
task_kwargs = {
"description": formatted_description,
"expected_output": config["expected_output"],
"agent": agent
}
if output_pydantic:
task_kwargs["output_pydantic"] = output_pydantic
return Task(**task_kwargs)
def create_synthesis_task(self, query: str, filtered_context: Dict[str, Any], agent) -> Task:
config = self.config_loader.get_task_config("synthesis_task")
return Task(
description=config["description"].format(
query=query,
filtered_context=json.dumps(filtered_context, indent=2)
),
expected_output=config["expected_output"],
agent=agent
)
File diff suppressed because it is too large Load Diff