Merge pull request #114 from Madhuvod/r1-agent-ollama
Added new demo R1 + RAG locally
This commit is contained in:
commit
6c3dd676b5
3 changed files with 660 additions and 0 deletions
87
ai_agent_tutorials/ai_knowledge_companion_r1_agent/README.md
Normal file
87
ai_agent_tutorials/ai_knowledge_companion_r1_agent/README.md
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
# Deepseek r1 Knowledge Agent 🤔
|
||||
|
||||
A versatile knowledge companion built with Deepseek r1 (via Ollama), Gemini for embeddings, Qdrant for vector storage, and Agno for agent orchestration. This application features dual-mode operation - a simple chat mode using local Deepseek r1 and an advanced RAG mode with document processing and web search capabilities.
|
||||
|
||||
## Features
|
||||
|
||||
- **Dual Operation Modes**
|
||||
- Simple Chat Mode: Direct interaction with Deepseek r1 locally
|
||||
- RAG Mode: Enhanced responses with document context and web search
|
||||
|
||||
- **Document Processing** (RAG Mode)
|
||||
- PDF document upload and processing
|
||||
- Web page content extraction
|
||||
- Automatic text chunking
|
||||
- Vector storage in Qdrant cloud
|
||||
|
||||
- **Intelligent Querying** (RAG Mode)
|
||||
- Query rewriting using Gemini
|
||||
- RAG-based document retrieval
|
||||
- Similarity search with threshold filtering
|
||||
- Automatic fallback to web search
|
||||
- Source attribution for answers
|
||||
|
||||
- **Advanced Capabilities**
|
||||
- Exa AI web search integration
|
||||
- Custom domain filtering for web search
|
||||
- Context-aware response generation
|
||||
- Chat history management
|
||||
- Thinking process visualization
|
||||
|
||||
- **Model Specific Features**
|
||||
- Flexible model selection:
|
||||
- Deepseek r1 1.5b (lighter, suitable for most laptops)
|
||||
- Deepseek r1 7b (more capable, requires better hardware)
|
||||
- Gemini Embedding model for vector embeddings
|
||||
- Agno Agent framework for orchestration
|
||||
- Streamlit-based interactive interface
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### 1. Ollama Setup
|
||||
1. Install [Ollama](https://ollama.ai)
|
||||
2. Pull the Deepseek r1 model(s):
|
||||
```bash
|
||||
# For the lighter model
|
||||
ollama pull deepseek-r1:1.5b
|
||||
|
||||
# For the more capable model (if your hardware supports it)
|
||||
ollama pull deepseek-r1:7b
|
||||
```
|
||||
|
||||
### 2. Google API Key (for RAG Mode)
|
||||
1. Go to [Google AI Studio](https://aistudio.google.com/apikey)
|
||||
2. Sign up or log in to your account
|
||||
3. Create a new API key
|
||||
|
||||
### 3. Qdrant Cloud Setup (for RAG Mode)
|
||||
1. Visit [Qdrant Cloud](https://cloud.qdrant.io/)
|
||||
2. Create an account or sign in
|
||||
3. Create a new cluster
|
||||
4. Get your credentials:
|
||||
- Qdrant API Key: Found in API Keys section
|
||||
- Qdrant URL: Your cluster URL (format: `https://xxx-xxx.cloud.qdrant.io`)
|
||||
|
||||
### 4. Exa AI API Key (Optional)
|
||||
1. Visit [Exa AI](https://exa.ai)
|
||||
2. Sign up for an account
|
||||
3. Generate an API key for web search capabilities
|
||||
|
||||
## How to Run
|
||||
|
||||
1. Clone the repository:
|
||||
```bash
|
||||
git clone https://github.com/Shubhamsaboo/awesome-llm-apps.git
|
||||
cd ai_agent_tutorials/ai_knowledge_companion_r1_agent
|
||||
```
|
||||
|
||||
2. Install dependencies:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
3. Run the application:
|
||||
```bash
|
||||
streamlit run ai_knowledge_r1_agent.py
|
||||
```
|
||||
|
||||
|
|
@ -0,0 +1,566 @@
|
|||
import os
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from typing import List
|
||||
|
||||
import streamlit as st
|
||||
import google.generativeai as genai
|
||||
import bs4
|
||||
from agno.agent import Agent
|
||||
from agno.models.google import Gemini
|
||||
from agno.models.ollama import Ollama
|
||||
from langchain_community.document_loaders import PyPDFLoader, WebBaseLoader
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import Distance, VectorParams
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from agno.tools.exa import ExaTools
|
||||
|
||||
|
||||
class GeminiEmbedder(Embeddings):
|
||||
def __init__(self, model_name="models/text-embedding-004"):
|
||||
genai.configure(api_key=st.session_state.google_api_key)
|
||||
self.model = model_name
|
||||
|
||||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||||
return [self.embed_query(text) for text in texts]
|
||||
|
||||
def embed_query(self, text: str) -> List[float]:
|
||||
response = genai.embed_content(
|
||||
model=self.model,
|
||||
content=text,
|
||||
task_type="retrieval_document"
|
||||
)
|
||||
return response['embedding']
|
||||
|
||||
|
||||
# Constants
|
||||
COLLECTION_NAME = "deepseek-r1-agno"
|
||||
|
||||
|
||||
# Streamlit App Initialization
|
||||
st.title("🤔 Deepseek r1 Knowledge Agent ")
|
||||
|
||||
# Session State Initialization
|
||||
if 'google_api_key' not in st.session_state:
|
||||
st.session_state.google_api_key = ""
|
||||
if 'qdrant_api_key' not in st.session_state:
|
||||
st.session_state.qdrant_api_key = ""
|
||||
if 'qdrant_url' not in st.session_state:
|
||||
st.session_state.qdrant_url = ""
|
||||
if 'model_version' not in st.session_state:
|
||||
st.session_state.model_version = "deepseek-r1:1.5b" # Default to lighter model
|
||||
if 'vector_store' not in st.session_state:
|
||||
st.session_state.vector_store = None
|
||||
if 'processed_documents' not in st.session_state:
|
||||
st.session_state.processed_documents = []
|
||||
if 'history' not in st.session_state:
|
||||
st.session_state.history = []
|
||||
if 'exa_api_key' not in st.session_state:
|
||||
st.session_state.exa_api_key = ""
|
||||
if 'use_web_search' not in st.session_state:
|
||||
st.session_state.use_web_search = False
|
||||
if 'force_web_search' not in st.session_state:
|
||||
st.session_state.force_web_search = False
|
||||
if 'similarity_threshold' not in st.session_state:
|
||||
st.session_state.similarity_threshold = 0.7
|
||||
if 'rag_enabled' not in st.session_state:
|
||||
st.session_state.rag_enabled = True # RAG is enabled by default
|
||||
|
||||
|
||||
# Sidebar Configuration
|
||||
st.sidebar.header("🤖 Agent Configuration")
|
||||
|
||||
# Model Selection
|
||||
st.sidebar.header("📦 Model Selection")
|
||||
model_help = """
|
||||
- 1.5b: Lighter model, suitable for most laptops
|
||||
- 7b: More capable but requires better GPU/RAM
|
||||
|
||||
Choose based on your hardware capabilities.
|
||||
"""
|
||||
st.session_state.model_version = st.sidebar.radio(
|
||||
"Select Model Version",
|
||||
options=["deepseek-r1:1.5b", "deepseek-r1:7b"],
|
||||
help=model_help
|
||||
)
|
||||
st.sidebar.info("Run ollama pull deepseek-r1:7b or deepseek-r1:1.5b respectively")
|
||||
|
||||
# RAG Mode Toggle
|
||||
st.sidebar.header("🔍 RAG Configuration")
|
||||
st.session_state.rag_enabled = st.sidebar.toggle("Enable RAG Mode", value=st.session_state.rag_enabled)
|
||||
|
||||
# Clear Chat Button
|
||||
if st.sidebar.button("🗑️ Clear Chat History"):
|
||||
st.session_state.history = []
|
||||
st.rerun()
|
||||
|
||||
# Show API Configuration only if RAG is enabled
|
||||
if st.session_state.rag_enabled:
|
||||
st.sidebar.header("🔑 API Configuration")
|
||||
google_api_key = st.sidebar.text_input("Google API Key", type="password", value=st.session_state.google_api_key)
|
||||
qdrant_api_key = st.sidebar.text_input("Qdrant API Key", type="password", value=st.session_state.qdrant_api_key)
|
||||
qdrant_url = st.sidebar.text_input("Qdrant URL",
|
||||
placeholder="https://your-cluster.cloud.qdrant.io:6333",
|
||||
value=st.session_state.qdrant_url)
|
||||
|
||||
# Update session state
|
||||
st.session_state.google_api_key = google_api_key
|
||||
st.session_state.qdrant_api_key = qdrant_api_key
|
||||
st.session_state.qdrant_url = qdrant_url
|
||||
|
||||
# Search Configuration (only shown in RAG mode)
|
||||
st.sidebar.header("🎯 Search Configuration")
|
||||
st.session_state.similarity_threshold = st.sidebar.slider(
|
||||
"Document Similarity Threshold",
|
||||
min_value=0.0,
|
||||
max_value=1.0,
|
||||
value=0.7,
|
||||
help="Lower values will return more documents but might be less relevant. Higher values are more strict."
|
||||
)
|
||||
|
||||
# Add in the sidebar configuration section, after the existing API inputs
|
||||
|
||||
st.sidebar.header("🌐 Web Search Configuration")
|
||||
st.session_state.use_web_search = st.sidebar.checkbox("Enable Web Search Fallback", value=st.session_state.use_web_search)
|
||||
|
||||
if st.session_state.use_web_search:
|
||||
exa_api_key = st.sidebar.text_input(
|
||||
"Exa AI API Key",
|
||||
type="password",
|
||||
value=st.session_state.exa_api_key,
|
||||
help="Required for web search fallback when no relevant documents are found"
|
||||
)
|
||||
st.session_state.exa_api_key = exa_api_key
|
||||
|
||||
# Optional domain filtering
|
||||
default_domains = ["arxiv.org", "wikipedia.org", "github.com", "medium.com"]
|
||||
custom_domains = st.sidebar.text_input(
|
||||
"Custom domains (comma-separated)",
|
||||
value=",".join(default_domains),
|
||||
help="Enter domains to search from, e.g.: arxiv.org,wikipedia.org"
|
||||
)
|
||||
search_domains = [d.strip() for d in custom_domains.split(",") if d.strip()]
|
||||
|
||||
# Search Configuration moved inside RAG mode check
|
||||
|
||||
|
||||
# Utility Functions
|
||||
def init_qdrant():
|
||||
"""Initialize Qdrant client with configured settings."""
|
||||
if not all([st.session_state.qdrant_api_key, st.session_state.qdrant_url]):
|
||||
return None
|
||||
try:
|
||||
return QdrantClient(
|
||||
url=st.session_state.qdrant_url,
|
||||
api_key=st.session_state.qdrant_api_key,
|
||||
timeout=60
|
||||
)
|
||||
except Exception as e:
|
||||
st.error(f"🔴 Qdrant connection failed: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
# Document Processing Functions
|
||||
def process_pdf(file) -> List:
|
||||
"""Process PDF file and add source metadata."""
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
|
||||
tmp_file.write(file.getvalue())
|
||||
loader = PyPDFLoader(tmp_file.name)
|
||||
documents = loader.load()
|
||||
|
||||
# Add source metadata
|
||||
for doc in documents:
|
||||
doc.metadata.update({
|
||||
"source_type": "pdf",
|
||||
"file_name": file.name,
|
||||
"timestamp": datetime.now().isoformat()
|
||||
})
|
||||
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200
|
||||
)
|
||||
return text_splitter.split_documents(documents)
|
||||
except Exception as e:
|
||||
st.error(f"📄 PDF processing error: {str(e)}")
|
||||
return []
|
||||
|
||||
|
||||
def process_web(url: str) -> List:
|
||||
"""Process web URL and add source metadata."""
|
||||
try:
|
||||
loader = WebBaseLoader(
|
||||
web_paths=(url,),
|
||||
bs_kwargs=dict(
|
||||
parse_only=bs4.SoupStrainer(
|
||||
class_=("post-content", "post-title", "post-header", "content", "main")
|
||||
)
|
||||
)
|
||||
)
|
||||
documents = loader.load()
|
||||
|
||||
# Add source metadata
|
||||
for doc in documents:
|
||||
doc.metadata.update({
|
||||
"source_type": "url",
|
||||
"url": url,
|
||||
"timestamp": datetime.now().isoformat()
|
||||
})
|
||||
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200
|
||||
)
|
||||
return text_splitter.split_documents(documents)
|
||||
except Exception as e:
|
||||
st.error(f"🌐 Web processing error: {str(e)}")
|
||||
return []
|
||||
|
||||
|
||||
# Vector Store Management
|
||||
def create_vector_store(client, texts):
|
||||
"""Create and initialize vector store with documents."""
|
||||
try:
|
||||
# Create collection if needed
|
||||
try:
|
||||
client.create_collection(
|
||||
collection_name=COLLECTION_NAME,
|
||||
vectors_config=VectorParams(
|
||||
size=768, # Gemini embedding-004 dimension
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
st.success(f"📚 Created new collection: {COLLECTION_NAME}")
|
||||
except Exception as e:
|
||||
if "already exists" not in str(e).lower():
|
||||
raise e
|
||||
|
||||
# Initialize vector store
|
||||
vector_store = QdrantVectorStore(
|
||||
client=client,
|
||||
collection_name=COLLECTION_NAME,
|
||||
embedding=GeminiEmbedder()
|
||||
)
|
||||
|
||||
# Add documents
|
||||
with st.spinner('📤 Uploading documents to Qdrant...'):
|
||||
vector_store.add_documents(texts)
|
||||
st.success("✅ Documents stored successfully!")
|
||||
return vector_store
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"🔴 Vector store error: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
# Add this after the GeminiEmbedder class
|
||||
def get_query_rewriter_agent() -> Agent:
|
||||
"""Initialize a query rewriting agent."""
|
||||
return Agent(
|
||||
name="Query Rewriter",
|
||||
model=Gemini(id="gemini-exp-1206"),
|
||||
instructions="""You are an expert at reformulating questions to be more precise and detailed.
|
||||
1. Analyze the user's question
|
||||
2. Rewrite it to be more specific and search-friendly
|
||||
3. Expand any acronyms or technical terms
|
||||
4. Return ONLY the rewritten query without any additional text or explanations
|
||||
|
||||
Example 1:
|
||||
User: "What does it say about ML?"
|
||||
Output: "What are the key concepts, techniques, and applications of Machine Learning (ML) discussed in the context?"
|
||||
|
||||
Example 2:
|
||||
User: "Tell me about transformers"
|
||||
Output: "Explain the architecture, mechanisms, and applications of Transformer neural networks in natural language processing and deep learning"
|
||||
""",
|
||||
show_tool_calls=False,
|
||||
markdown=True,
|
||||
)
|
||||
|
||||
|
||||
def get_web_search_agent() -> Agent:
|
||||
"""Initialize a web search agent."""
|
||||
return Agent(
|
||||
name="Web Search Agent",
|
||||
model=Gemini(id="gemini-exp-1206", api_key=st.session_state.google_api_key),
|
||||
tools=[ExaTools(
|
||||
api_key=st.session_state.exa_api_key,
|
||||
include_domains=search_domains,
|
||||
num_results=5
|
||||
)],
|
||||
instructions="""You are a web search expert. Your task is to:
|
||||
1. Search the web for relevant information about the query
|
||||
2. Compile and summarize the most relevant information
|
||||
3. Include sources in your response
|
||||
""",
|
||||
show_tool_calls=True,
|
||||
markdown=True,
|
||||
)
|
||||
|
||||
|
||||
def get_rag_agent() -> Agent:
|
||||
"""Initialize the main RAG agent."""
|
||||
return Agent(
|
||||
name="DeepSeek RAG Agent",
|
||||
model=Ollama(id=st.session_state.model_version),
|
||||
instructions="""You are an Intelligent Agent specializing in providing accurate answers.
|
||||
|
||||
When asked a question:
|
||||
- Analyze the question and answer the question with what you know.
|
||||
|
||||
When given context from documents:
|
||||
- Focus on information from the provided documents
|
||||
- Be precise and cite specific details
|
||||
|
||||
When given web search results:
|
||||
- Clearly indicate that the information comes from web search
|
||||
- Synthesize the information clearly
|
||||
|
||||
Always maintain high accuracy and clarity in your responses.
|
||||
""",
|
||||
show_tool_calls=True,
|
||||
markdown=True,
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
def check_document_relevance(query: str, vector_store, threshold: float = 0.7) -> tuple[bool, List]:
|
||||
|
||||
if not vector_store:
|
||||
return False, []
|
||||
|
||||
retriever = vector_store.as_retriever(
|
||||
search_type="similarity_score_threshold",
|
||||
search_kwargs={"k": 5, "score_threshold": threshold}
|
||||
)
|
||||
docs = retriever.invoke(query)
|
||||
return bool(docs), docs
|
||||
|
||||
|
||||
# Main Application Flow
|
||||
|
||||
# Chat Interface
|
||||
# Create two columns for chat input and search toggle
|
||||
chat_col, toggle_col = st.columns([0.9, 0.1])
|
||||
|
||||
with chat_col:
|
||||
prompt = st.chat_input("Ask about your documents..." if st.session_state.rag_enabled else "Ask me anything...")
|
||||
|
||||
with toggle_col:
|
||||
st.session_state.force_web_search = st.toggle('🌐', help="Force web search")
|
||||
|
||||
# Check if RAG is enabled but no API key
|
||||
if st.session_state.rag_enabled and not st.session_state.google_api_key:
|
||||
st.error("Please enter your Google API Key to continue with RAG mode")
|
||||
st.stop()
|
||||
|
||||
# Initialize Qdrant and configure APIs if RAG is enabled
|
||||
if st.session_state.rag_enabled and st.session_state.google_api_key:
|
||||
os.environ["GOOGLE_API_KEY"] = st.session_state.google_api_key
|
||||
genai.configure(api_key=st.session_state.google_api_key)
|
||||
|
||||
qdrant_client = init_qdrant()
|
||||
|
||||
# File/URL Upload Section
|
||||
st.sidebar.header("📁 Data Upload")
|
||||
uploaded_file = st.sidebar.file_uploader("Upload PDF", type=["pdf"])
|
||||
web_url = st.sidebar.text_input("Or enter URL")
|
||||
|
||||
# Process documents
|
||||
if uploaded_file:
|
||||
file_name = uploaded_file.name
|
||||
if file_name not in st.session_state.processed_documents:
|
||||
with st.spinner('Processing PDF...'):
|
||||
texts = process_pdf(uploaded_file)
|
||||
if texts and qdrant_client:
|
||||
if st.session_state.vector_store:
|
||||
st.session_state.vector_store.add_documents(texts)
|
||||
else:
|
||||
st.session_state.vector_store = create_vector_store(qdrant_client, texts)
|
||||
st.session_state.processed_documents.append(file_name)
|
||||
st.success(f"✅ Added PDF: {file_name}")
|
||||
|
||||
if web_url:
|
||||
if web_url not in st.session_state.processed_documents:
|
||||
with st.spinner('Processing URL...'):
|
||||
texts = process_web(web_url)
|
||||
if texts and qdrant_client:
|
||||
if st.session_state.vector_store:
|
||||
st.session_state.vector_store.add_documents(texts)
|
||||
else:
|
||||
st.session_state.vector_store = create_vector_store(qdrant_client, texts)
|
||||
st.session_state.processed_documents.append(web_url)
|
||||
st.success(f"✅ Added URL: {web_url}")
|
||||
|
||||
# Display sources in sidebar
|
||||
if st.session_state.processed_documents:
|
||||
st.sidebar.header("📚 Processed Sources")
|
||||
for source in st.session_state.processed_documents:
|
||||
if source.endswith('.pdf'):
|
||||
st.sidebar.text(f"📄 {source}")
|
||||
else:
|
||||
st.sidebar.text(f"🌐 {source}")
|
||||
|
||||
if prompt:
|
||||
# Add user message to history
|
||||
st.session_state.history.append({"role": "user", "content": prompt})
|
||||
with st.chat_message("user"):
|
||||
st.write(prompt)
|
||||
|
||||
if st.session_state.rag_enabled:
|
||||
# Existing RAG flow remains unchanged
|
||||
with st.spinner("🤔 Reformulating query..."):
|
||||
try:
|
||||
query_rewriter = get_query_rewriter_agent()
|
||||
rewritten_query = query_rewriter.run(prompt).content
|
||||
|
||||
with st.expander("🔄 See rewritten query"):
|
||||
st.write(f"Original: {prompt}")
|
||||
st.write(f"Rewritten: {rewritten_query}")
|
||||
except Exception as e:
|
||||
st.error(f"❌ Error rewriting query: {str(e)}")
|
||||
rewritten_query = prompt
|
||||
|
||||
# Step 2: Choose search strategy based on force_web_search toggle
|
||||
context = ""
|
||||
docs = []
|
||||
if not st.session_state.force_web_search and st.session_state.vector_store:
|
||||
# Try document search first
|
||||
retriever = st.session_state.vector_store.as_retriever(
|
||||
search_type="similarity_score_threshold",
|
||||
search_kwargs={
|
||||
"k": 5,
|
||||
"score_threshold": st.session_state.similarity_threshold
|
||||
}
|
||||
)
|
||||
docs = retriever.invoke(rewritten_query)
|
||||
if docs:
|
||||
context = "\n\n".join([d.page_content for d in docs])
|
||||
st.info(f"📊 Found {len(docs)} relevant documents (similarity > {st.session_state.similarity_threshold})")
|
||||
elif st.session_state.use_web_search:
|
||||
st.info("🔄 No relevant documents found in database, falling back to web search...")
|
||||
|
||||
# Step 3: Use web search if:
|
||||
# 1. Web search is forced ON via toggle, or
|
||||
# 2. No relevant documents found AND web search is enabled in settings
|
||||
if (st.session_state.force_web_search or not context) and st.session_state.use_web_search and st.session_state.exa_api_key:
|
||||
with st.spinner("🔍 Searching the web..."):
|
||||
try:
|
||||
web_search_agent = get_web_search_agent()
|
||||
web_results = web_search_agent.run(rewritten_query).content
|
||||
if web_results:
|
||||
context = f"Web Search Results:\n{web_results}"
|
||||
if st.session_state.force_web_search:
|
||||
st.info("ℹ️ Using web search as requested via toggle.")
|
||||
else:
|
||||
st.info("ℹ️ Using web search as fallback since no relevant documents were found.")
|
||||
except Exception as e:
|
||||
st.error(f"❌ Web search error: {str(e)}")
|
||||
|
||||
# Step 4: Generate response using the RAG agent
|
||||
with st.spinner("🤖 Thinking..."):
|
||||
try:
|
||||
rag_agent = get_rag_agent()
|
||||
|
||||
if context:
|
||||
full_prompt = f"""Context: {context}
|
||||
|
||||
Original Question: {prompt}
|
||||
Rewritten Question: {rewritten_query}
|
||||
|
||||
Please provide a comprehensive answer based on the available information."""
|
||||
else:
|
||||
full_prompt = f"Original Question: {prompt}\nRewritten Question: {rewritten_query}"
|
||||
st.info("ℹ️ No relevant information found in documents or web search.")
|
||||
|
||||
response = rag_agent.run(full_prompt)
|
||||
|
||||
# Add assistant response to history
|
||||
st.session_state.history.append({
|
||||
"role": "assistant",
|
||||
"content": response.content
|
||||
})
|
||||
|
||||
# Display assistant response
|
||||
with st.chat_message("assistant"):
|
||||
st.write(response.content)
|
||||
|
||||
# Show sources if available
|
||||
if not st.session_state.force_web_search and 'docs' in locals() and docs:
|
||||
with st.expander("🔍 See document sources"):
|
||||
for i, doc in enumerate(docs, 1):
|
||||
source_type = doc.metadata.get("source_type", "unknown")
|
||||
source_icon = "📄" if source_type == "pdf" else "🌐"
|
||||
source_name = doc.metadata.get("file_name" if source_type == "pdf" else "url", "unknown")
|
||||
st.write(f"{source_icon} Source {i} from {source_name}:")
|
||||
st.write(f"{doc.page_content[:200]}...")
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"❌ Error generating response: {str(e)}")
|
||||
|
||||
else:
|
||||
# Simple mode without RAG
|
||||
with st.spinner("🤖 Thinking..."):
|
||||
try:
|
||||
rag_agent = get_rag_agent()
|
||||
web_search_agent = get_web_search_agent() if st.session_state.use_web_search else None
|
||||
|
||||
# Handle web search if forced or enabled
|
||||
context = ""
|
||||
if st.session_state.force_web_search and web_search_agent:
|
||||
with st.spinner("🔍 Searching the web..."):
|
||||
try:
|
||||
web_results = web_search_agent.run(prompt).content
|
||||
if web_results:
|
||||
context = f"Web Search Results:\n{web_results}"
|
||||
st.info("ℹ️ Using web search as requested.")
|
||||
except Exception as e:
|
||||
st.error(f"❌ Web search error: {str(e)}")
|
||||
|
||||
# Generate response
|
||||
if context:
|
||||
full_prompt = f"""Context: {context}
|
||||
|
||||
Question: {prompt}
|
||||
|
||||
Please provide a comprehensive answer based on the available information."""
|
||||
else:
|
||||
full_prompt = prompt
|
||||
|
||||
response = rag_agent.run(full_prompt)
|
||||
response_content = response.content
|
||||
|
||||
# Extract thinking process and final response
|
||||
import re
|
||||
think_pattern = r'<think>(.*?)</think>'
|
||||
think_match = re.search(think_pattern, response_content, re.DOTALL)
|
||||
|
||||
if think_match:
|
||||
thinking_process = think_match.group(1).strip()
|
||||
final_response = re.sub(think_pattern, '', response_content, flags=re.DOTALL).strip()
|
||||
else:
|
||||
thinking_process = None
|
||||
final_response = response_content
|
||||
|
||||
# Add assistant response to history (only the final response)
|
||||
st.session_state.history.append({
|
||||
"role": "assistant",
|
||||
"content": final_response
|
||||
})
|
||||
|
||||
# Display assistant response
|
||||
with st.chat_message("assistant"):
|
||||
if thinking_process:
|
||||
with st.expander("🤔 See thinking process"):
|
||||
st.markdown(thinking_process)
|
||||
st.markdown(final_response)
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"❌ Error generating response: {str(e)}")
|
||||
|
||||
else:
|
||||
st.warning("You can directly talk to r1 locally! Toggle the RAG mode to upload documents!")
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
agno
|
||||
exa==0.5.26
|
||||
qdrant-client==1.12.1
|
||||
langchain-qdrant==0.2.0
|
||||
langchain-community==0.3.13
|
||||
streamlit==1.41.1
|
||||
ollama
|
||||
Loading…
Reference in a new issue