rag document, url working
This commit is contained in:
Madhu
parent
7aa9a761ff
commit
40ac9f865d
2 changed files with 119 additions and 50 deletions
|
|
@ -1,9 +1,11 @@
|
|||
import os
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from typing import List
|
||||
|
||||
import streamlit as st
|
||||
import google.generativeai as genai
|
||||
import tempfile
|
||||
import bs4
|
||||
from typing import List
|
||||
from agno.agent import Agent
|
||||
from agno.models.google import Gemini
|
||||
from langchain_community.document_loaders import PyPDFLoader, WebBaseLoader
|
||||
|
|
@ -14,10 +16,10 @@ from qdrant_client.models import Distance, VectorParams
|
|||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
|
||||
# Custom Gemini Embedder Class
|
||||
# Custom Classes
|
||||
class GeminiEmbedder(Embeddings):
|
||||
def __init__(self, model_name="models/embedding-004"):
|
||||
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
|
||||
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]]:
|
||||
|
|
@ -31,37 +33,74 @@ class GeminiEmbedder(Embeddings):
|
|||
)
|
||||
return response['embedding']
|
||||
|
||||
# Initialize Streamlit App
|
||||
|
||||
# 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")
|
||||
qdrant_api_key = st.sidebar.text_input("Qdrant API Key", type="password")
|
||||
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")
|
||||
placeholder="https://your-cluster.cloud.qdrant.io:6333",
|
||||
value=st.session_state.qdrant_url)
|
||||
|
||||
# Initialize Qdrant Client
|
||||
# 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():
|
||||
if not all([qdrant_api_key, qdrant_url]):
|
||||
"""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=qdrant_url,
|
||||
api_key=qdrant_api_key,
|
||||
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):
|
||||
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
|
||||
|
|
@ -71,7 +110,9 @@ def process_pdf(file):
|
|||
st.error(f"📄 PDF processing error: {str(e)}")
|
||||
return []
|
||||
|
||||
def process_web(url):
|
||||
|
||||
def process_web(url: str) -> List:
|
||||
"""Process web URL and add source metadata."""
|
||||
try:
|
||||
loader = WebBaseLoader(
|
||||
web_paths=(url,),
|
||||
|
|
@ -82,6 +123,15 @@ def process_web(url):
|
|||
)
|
||||
)
|
||||
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
|
||||
|
|
@ -91,10 +141,10 @@ def process_web(url):
|
|||
st.error(f"🌐 Web processing error: {str(e)}")
|
||||
return []
|
||||
|
||||
# Vector Store Management
|
||||
COLLECTION_NAME = "agno_rag"
|
||||
|
||||
# Vector Store Management
|
||||
def create_vector_store(client, texts):
|
||||
"""Create and initialize vector store with documents."""
|
||||
try:
|
||||
# Create collection if needed
|
||||
try:
|
||||
|
|
@ -127,10 +177,11 @@ def create_vector_store(client, texts):
|
|||
st.error(f"🔴 Vector store error: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
# Main Application Flow
|
||||
if google_api_key:
|
||||
os.environ["GOOGLE_API_KEY"] = google_api_key
|
||||
genai.configure(api_key=google_api_key)
|
||||
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()
|
||||
|
||||
|
|
@ -140,15 +191,39 @@ if google_api_key:
|
|||
web_url = st.sidebar.text_input("Or enter URL")
|
||||
|
||||
# Process documents
|
||||
vector_store = None
|
||||
if uploaded_file:
|
||||
texts = process_pdf(uploaded_file)
|
||||
if texts and qdrant_client:
|
||||
vector_store = create_vector_store(qdrant_client, texts)
|
||||
elif web_url:
|
||||
texts = process_web(web_url)
|
||||
if texts and qdrant_client:
|
||||
vector_store = create_vector_store(qdrant_client, texts)
|
||||
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(
|
||||
|
|
@ -159,7 +234,7 @@ if google_api_key:
|
|||
markdown=True,
|
||||
)
|
||||
|
||||
# Initialize chat history
|
||||
# Chat Interface
|
||||
if 'history' not in st.session_state:
|
||||
st.session_state.history = []
|
||||
|
||||
|
|
@ -168,7 +243,7 @@ if google_api_key:
|
|||
with st.chat_message(msg["role"]):
|
||||
st.write(msg["content"])
|
||||
|
||||
# User input
|
||||
# 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})
|
||||
|
|
@ -177,8 +252,8 @@ if google_api_key:
|
|||
|
||||
# Retrieve relevant documents
|
||||
context = ""
|
||||
if vector_store:
|
||||
retriever = vector_store.as_retriever(
|
||||
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}
|
||||
)
|
||||
|
|
@ -200,10 +275,14 @@ if google_api_key:
|
|||
with st.chat_message("assistant"):
|
||||
st.write(response.content)
|
||||
|
||||
if vector_store and docs:
|
||||
if st.session_state.vector_store and docs:
|
||||
with st.expander("🔍 See sources"):
|
||||
for i, doc in enumerate(docs, 1):
|
||||
st.write(f"Source {i}: {doc.page_content[:200]}...")
|
||||
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)}")
|
||||
|
|
|
|||
|
|
@ -1,22 +1,12 @@
|
|||
from google import genai
|
||||
import google.generativeai as genai
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"), http_options={'api_version':'v1alpha'})
|
||||
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
|
||||
|
||||
import asyncio
|
||||
result = genai.embed_content(
|
||||
model="models/text-embedding-004",
|
||||
content="What is the meaning of life?")
|
||||
|
||||
config = {'thinking_config': {'include_thoughts': True}}
|
||||
|
||||
async def main():
|
||||
chat = client.aio.chats.create(
|
||||
model='gemini-2.0-flash-thinking-exp-01-21',
|
||||
config=config
|
||||
)
|
||||
response = await chat.send_message('Explain Deep Q Networks from first principles')
|
||||
print(response.text)
|
||||
response = await chat.send_message('What did you just say before this?')
|
||||
print(response.text)
|
||||
|
||||
asyncio.run(main())
|
||||
print(str(result['embedding']))
|
||||
Loading…
Reference in a new issue