main script working
This commit is contained in:
Madhu
parent
610c8adee8
commit
1eb3e4cff6
3 changed files with 546 additions and 0 deletions
|
|
@ -0,0 +1,546 @@
|
||||||
|
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 '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")
|
||||||
|
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="deepseek-r1:1.5b"),
|
||||||
|
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!")
|
||||||
Loading…
Reference in a new issue