New project with gemini thinking3
This commit is contained in:
Madhu
parent
38911478ce
commit
7aa9a761ff
1 changed files with 146 additions and 105 deletions
|
|
@ -1,171 +1,212 @@
|
|||
import os
|
||||
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 agno.tools.duckduckgo import DuckDuckGoTools
|
||||
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 agno.vectordb.pgvector import PgVector
|
||||
from agno.embedder.google import GeminiEmbedder
|
||||
import tempfile
|
||||
import bs4
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
# Streamlit App Title
|
||||
st.title("AI Agent with Agno and Gemini Thinking")
|
||||
|
||||
# Sidebar for API Key Input
|
||||
st.sidebar.header("Configuration")
|
||||
google_api_key = st.sidebar.text_input("Enter your Google API Key", type="password")
|
||||
qdrant_api_key = st.sidebar.text_input("Enter your Qdrant API Key", type="password")
|
||||
qdrant_url = st.sidebar.text_input("Enter your Qdrant URL", placeholder="https://your-qdrant-url.com")
|
||||
# Custom Gemini Embedder Class
|
||||
class GeminiEmbedder(Embeddings):
|
||||
def __init__(self, model_name="models/embedding-004"):
|
||||
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
|
||||
self.model = model_name
|
||||
|
||||
if google_api_key:
|
||||
os.environ["GOOGLE_API_KEY"] = google_api_key
|
||||
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']
|
||||
|
||||
# Initialize Streamlit App
|
||||
st.title("🤖 AI Agent with Gemini & Qdrant RAG")
|
||||
|
||||
# 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")
|
||||
qdrant_url = st.sidebar.text_input("Qdrant URL",
|
||||
placeholder="https://your-cluster.cloud.qdrant.io:6333")
|
||||
|
||||
# Initialize Qdrant Client
|
||||
def init_qdrant():
|
||||
if not qdrant_api_key or not qdrant_url:
|
||||
st.warning("Please provide Qdrant API Key and URL in the sidebar.")
|
||||
if not all([qdrant_api_key, qdrant_url]):
|
||||
return None
|
||||
try:
|
||||
client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key, timeout=60)
|
||||
client.get_collections() # Test connection
|
||||
return client
|
||||
return QdrantClient(
|
||||
url=qdrant_url,
|
||||
api_key=qdrant_api_key,
|
||||
timeout=60
|
||||
)
|
||||
except Exception as e:
|
||||
st.error(f"Failed to initialize Qdrant: {e}")
|
||||
st.error(f"🔴 Qdrant connection failed: {str(e)}")
|
||||
return None
|
||||
|
||||
qdrant_client = init_qdrant()
|
||||
|
||||
# File/URL Upload Section
|
||||
st.sidebar.header("Upload Data")
|
||||
uploaded_file = st.sidebar.file_uploader("Upload a document", type=["txt", "pdf", "jpg", "png"])
|
||||
web_url = st.sidebar.text_input("Enter a web URL")
|
||||
|
||||
# Document and Web URL Processing
|
||||
def process_document(file):
|
||||
# Document Processing Functions
|
||||
def process_pdf(file):
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
|
||||
tmp_file.write(file.getvalue())
|
||||
tmp_path = tmp_file.name
|
||||
|
||||
loader = PyPDFLoader(tmp_path)
|
||||
documents = loader.load()
|
||||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||||
texts = text_splitter.split_documents(documents)
|
||||
|
||||
os.unlink(tmp_path)
|
||||
return texts
|
||||
loader = PyPDFLoader(tmp_file.name)
|
||||
documents = loader.load()
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200
|
||||
)
|
||||
return text_splitter.split_documents(documents)
|
||||
except Exception as e:
|
||||
st.error(f"Error processing document: {e}")
|
||||
st.error(f"📄 PDF processing error: {str(e)}")
|
||||
return []
|
||||
|
||||
def process_web_url(url):
|
||||
def process_web(url):
|
||||
try:
|
||||
loader = WebBaseLoader(
|
||||
web_paths=(url,),
|
||||
bs_kwargs=dict(
|
||||
parse_only=bs4.SoupStrainer(
|
||||
class_=("post-content", "post-title", "post-header")
|
||||
class_=("post-content", "post-title", "post-header", "content", "main")
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
documents = loader.load()
|
||||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||||
texts = text_splitter.split_documents(documents)
|
||||
return texts
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200
|
||||
)
|
||||
return text_splitter.split_documents(documents)
|
||||
except Exception as e:
|
||||
st.error(f"Error processing web URL: {e}")
|
||||
st.error(f"🌐 Web processing error: {str(e)}")
|
||||
return []
|
||||
|
||||
# Create and Populate Qdrant Vector Store
|
||||
# Vector Store Management
|
||||
COLLECTION_NAME = "agno_rag"
|
||||
|
||||
def create_vector_store(texts):
|
||||
if not qdrant_client:
|
||||
return None
|
||||
def create_vector_store(client, texts):
|
||||
try:
|
||||
# Create collection if it doesn't exist
|
||||
# Create collection if needed
|
||||
try:
|
||||
qdrant_client.create_collection(
|
||||
client.create_collection(
|
||||
collection_name=COLLECTION_NAME,
|
||||
vectors_config=VectorParams(size=1024, distance=Distance.COSINE)
|
||||
vectors_config=VectorParams(
|
||||
size=768, # Gemini embedding-004 dimension
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
st.success(f"Created new collection: {COLLECTION_NAME}")
|
||||
st.success(f"📚 Created new collection: {COLLECTION_NAME}")
|
||||
except Exception as e:
|
||||
if "already exists" not in str(e).lower():
|
||||
raise e
|
||||
|
||||
# Initialize QdrantVectorStore
|
||||
|
||||
# Initialize vector store
|
||||
vector_store = QdrantVectorStore(
|
||||
client=qdrant_client,
|
||||
client=client,
|
||||
collection_name=COLLECTION_NAME,
|
||||
embedding=GeminiEmbedder(dimensions=1024) # Add embedding model if needed
|
||||
embedding=GeminiEmbedder()
|
||||
)
|
||||
|
||||
# Add documents to the vector store
|
||||
with st.spinner('Storing documents in Qdrant...'):
|
||||
|
||||
# Add documents
|
||||
with st.spinner('📤 Uploading documents to Qdrant...'):
|
||||
vector_store.add_documents(texts)
|
||||
st.success("Documents successfully stored in Qdrant!")
|
||||
|
||||
return vector_store
|
||||
st.success("✅ Documents stored successfully!")
|
||||
return vector_store
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"Error creating vector store: {e}")
|
||||
st.error(f"🔴 Vector store error: {str(e)}")
|
||||
return None
|
||||
|
||||
# Process Uploaded File or Web URL
|
||||
if uploaded_file:
|
||||
texts = process_document(uploaded_file)
|
||||
if texts:
|
||||
vector_store = create_vector_store(texts)
|
||||
elif web_url:
|
||||
texts = process_web_url(web_url)
|
||||
if texts:
|
||||
vector_store = create_vector_store(texts)
|
||||
# Main Application Flow
|
||||
if google_api_key:
|
||||
os.environ["GOOGLE_API_KEY"] = google_api_key
|
||||
genai.configure(api_key=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
|
||||
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)
|
||||
|
||||
# Initialize the Agent
|
||||
if google_api_key and qdrant_client:
|
||||
thinking_agent = Agent(
|
||||
name="Thinking Agent",
|
||||
role="Think about the problem",
|
||||
model=Gemini(id="gemini-2.0-flash-exp", api_key=google_api_key),
|
||||
instructions="Given the problem, think about it and provide a detailed explanation",
|
||||
# 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,
|
||||
)
|
||||
|
||||
|
||||
# Display chat history if it exists
|
||||
if 'chat_history' not in st.session_state:
|
||||
st.session_state.chat_history = []
|
||||
# Initialize chat history
|
||||
if 'history' not in st.session_state:
|
||||
st.session_state.history = []
|
||||
|
||||
for message in st.session_state.chat_history:
|
||||
with st.chat_message(message["role"]):
|
||||
st.write(message["content"])
|
||||
# Display chat messages
|
||||
for msg in st.session_state.history:
|
||||
with st.chat_message(msg["role"]):
|
||||
st.write(msg["content"])
|
||||
|
||||
# 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)
|
||||
|
||||
# Chat input using Streamlit's chat_input for better UX
|
||||
user_input = st.chat_input("Ask a question or describe a problem you'd like me to think about...")
|
||||
|
||||
if user_input:
|
||||
# Query the Qdrant vector store for relevant documents
|
||||
if 'vector_store' in locals():
|
||||
# Retrieve relevant documents
|
||||
context = ""
|
||||
if vector_store:
|
||||
retriever = vector_store.as_retriever(
|
||||
search_type="similarity_score_threshold",
|
||||
search_kwargs={"k": 5, "score_threshold": 0.7}
|
||||
)
|
||||
relevant_docs = retriever.get_relevant_documents(user_input)
|
||||
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 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]}...")
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"❌ Error generating response: {str(e)}")
|
||||
|
||||
if relevant_docs:
|
||||
st.write("Relevant Documents:")
|
||||
for doc in relevant_docs:
|
||||
st.write(doc.page_content[:200] + "...")
|
||||
|
||||
# Process the user's input with the agent
|
||||
response = thinking_agent.run(user_input)
|
||||
st.write("Agent's Response:")
|
||||
st.write(response.content)
|
||||
else:
|
||||
st.warning("Please enter your Google API Key and Qdrant credentials in the sidebar to proceed.")
|
||||
st.warning("⚠️ Please enter your Google API Key to continue")
|
||||
Loading…
Reference in a new issue