291 lines
No EOL
11 KiB
Python
291 lines
No EOL
11 KiB
Python
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 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
|
|
|
|
|
|
# Custom Classes
|
|
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 = "gemini-rag-agno"
|
|
|
|
|
|
# Streamlit App Initialization
|
|
st.title("🤖 AI Agent with Gemini & Qdrant RAG")
|
|
|
|
# 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 = []
|
|
|
|
|
|
# Sidebar Configuration
|
|
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
|
|
|
|
|
|
# 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
|
|
|
|
|
|
# Main Application Flow
|
|
if 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}")
|
|
|
|
# Initialize Agent
|
|
agent = Agent(
|
|
name="Gemini RAG Agent",
|
|
model=Gemini(id="gemini-2.0-flash-exp"),
|
|
instructions="You are AGI. You are elite speicialist in all fields and an expert in all fields. Answer user's questions clearly, if any document is added, Use retrieved documents to answer questions accurately",
|
|
show_tool_calls=True,
|
|
markdown=True,
|
|
)
|
|
|
|
# Chat Interface
|
|
if 'history' not in st.session_state:
|
|
st.session_state.history = []
|
|
|
|
# Display chat messages
|
|
for msg in st.session_state.history:
|
|
with st.chat_message(msg["role"]):
|
|
st.write(msg["content"])
|
|
|
|
# Handle user input
|
|
if prompt := st.chat_input("Ask about your documents..."):
|
|
# Add user message to history
|
|
st.session_state.history.append({"role": "user", "content": prompt})
|
|
with st.chat_message("user"):
|
|
st.write(prompt)
|
|
|
|
# Retrieve relevant documents
|
|
context = ""
|
|
if st.session_state.vector_store:
|
|
retriever = st.session_state.vector_store.as_retriever(
|
|
search_type="similarity_score_threshold",
|
|
search_kwargs={"k": 5, "score_threshold": 0.7}
|
|
)
|
|
docs = retriever.invoke(prompt)
|
|
context = "\n\n".join([d.page_content for d in docs])
|
|
|
|
# Generate response
|
|
with st.spinner("🤖 Thinking..."):
|
|
try:
|
|
full_prompt = f"Context: {context}\n\nQuestion: {prompt}"
|
|
response = agent.run(full_prompt)
|
|
|
|
# Add assistant response to history
|
|
st.session_state.history.append({
|
|
"role": "assistant",
|
|
"content": response.content
|
|
})
|
|
|
|
with st.chat_message("assistant"):
|
|
st.write(response.content)
|
|
|
|
if st.session_state.vector_store and docs:
|
|
with st.expander("🔍 See 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:
|
|
st.warning("⚠️ Please enter your Google API Key to continue") |