complete working code
This commit is contained in:
Madhu
parent
7e8375831a
commit
93331106d7
1 changed files with 99 additions and 136 deletions
|
|
@ -12,10 +12,14 @@ from langchain import hub
|
|||
import tempfile
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
from langchain_community.tools import DuckDuckGoSearchRun
|
||||
from typing import TypedDict, List
|
||||
from langchain_core.language_models import BaseLanguageModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
||||
from time import sleep
|
||||
from tenacity import retry, wait_exponential, stop_after_attempt
|
||||
|
||||
|
||||
def init_session_state():
|
||||
"""Initialize session state variables."""
|
||||
if 'api_keys_submitted' not in st.session_state:
|
||||
st.session_state.api_keys_submitted = False
|
||||
if 'chat_history' not in st.session_state:
|
||||
|
|
@ -28,11 +32,9 @@ def init_session_state():
|
|||
st.session_state.qdrant_url = ""
|
||||
|
||||
def sidebar_api_form():
|
||||
"""Render API credentials form in sidebar."""
|
||||
with st.sidebar:
|
||||
st.header("API Credentials")
|
||||
|
||||
# Show current status
|
||||
if st.session_state.api_keys_submitted:
|
||||
st.success("API credentials verified")
|
||||
if st.button("Reset Credentials"):
|
||||
|
|
@ -40,32 +42,18 @@ def sidebar_api_form():
|
|||
st.rerun()
|
||||
return True
|
||||
|
||||
# Show API form
|
||||
with st.form("api_credentials"):
|
||||
cohere_key = st.text_input("Cohere API Key", type="password")
|
||||
qdrant_key = st.text_input(
|
||||
"Qdrant API Key",
|
||||
type="password",
|
||||
help="Enter your Qdrant API key"
|
||||
)
|
||||
qdrant_url = st.text_input(
|
||||
"Qdrant URL",
|
||||
placeholder="https://xyz-example.eu-central.aws.cloud.qdrant.io:6333",
|
||||
help="Enter your Qdrant instance URL"
|
||||
)
|
||||
qdrant_key = st.text_input("Qdrant API Key", type="password", help="Enter your Qdrant API key")
|
||||
qdrant_url = st.text_input("Qdrant URL",
|
||||
placeholder="https://xyz-example.eu-central.aws.cloud.qdrant.io:6333",
|
||||
help="Enter your Qdrant instance URL")
|
||||
|
||||
if st.form_submit_button("Submit Credentials"):
|
||||
try:
|
||||
# First validate the credentials before saving to session state
|
||||
client = QdrantClient(
|
||||
url=qdrant_url,
|
||||
api_key=qdrant_key,
|
||||
timeout=60
|
||||
)
|
||||
# Test connection
|
||||
client = QdrantClient(url=qdrant_url, api_key=qdrant_key, timeout=60)
|
||||
client.get_collections()
|
||||
|
||||
# Only save to session state after successful validation
|
||||
st.session_state.cohere_api_key = cohere_key
|
||||
st.session_state.qdrant_api_key = qdrant_key
|
||||
st.session_state.qdrant_url = qdrant_url
|
||||
|
|
@ -78,59 +66,43 @@ def sidebar_api_form():
|
|||
return False
|
||||
|
||||
def init_qdrant() -> QdrantClient:
|
||||
"""Initialize Qdrant vector database."""
|
||||
if not st.session_state.get("qdrant_api_key"):
|
||||
raise ValueError("Qdrant API key not provided")
|
||||
if not st.session_state.get("qdrant_url"):
|
||||
raise ValueError("Qdrant URL not provided")
|
||||
|
||||
return QdrantClient(
|
||||
url=st.session_state.qdrant_url,
|
||||
api_key=st.session_state.qdrant_api_key,
|
||||
timeout=60
|
||||
)
|
||||
return QdrantClient(url=st.session_state.qdrant_url,
|
||||
api_key=st.session_state.qdrant_api_key,
|
||||
timeout=60)
|
||||
|
||||
# Initialize session state
|
||||
init_session_state()
|
||||
|
||||
# Main application logic
|
||||
if not sidebar_api_form():
|
||||
st.info("Please enter your API credentials in the sidebar to continue.")
|
||||
st.stop()
|
||||
|
||||
# Initialize services with verified credentials
|
||||
embedding = CohereEmbeddings(
|
||||
model="embed-english-v3.0",
|
||||
cohere_api_key=st.session_state.cohere_api_key
|
||||
)
|
||||
embedding = CohereEmbeddings(model="embed-english-v3.0",
|
||||
cohere_api_key=st.session_state.cohere_api_key)
|
||||
|
||||
chat_model = ChatCohere(
|
||||
model="command-r7b-12-2024",
|
||||
temperature=0.1,
|
||||
max_tokens=512,
|
||||
verbose=True,
|
||||
cohere_api_key=st.session_state.cohere_api_key
|
||||
)
|
||||
chat_model = ChatCohere(model="command-r7b-12-2024",
|
||||
temperature=0.1,
|
||||
max_tokens=512,
|
||||
verbose=True,
|
||||
cohere_api_key=st.session_state.cohere_api_key)
|
||||
|
||||
client = init_qdrant()
|
||||
|
||||
#document preprocessing
|
||||
|
||||
def process_document(file):
|
||||
"""Process uploaded PDF document using a temporary file."""
|
||||
try:
|
||||
# Create a temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
|
||||
tmp_file.write(file.getvalue())
|
||||
tmp_path = tmp_file.name
|
||||
|
||||
# Process the temporary file
|
||||
loader = PyPDFLoader(tmp_path)
|
||||
documents = loader.load()
|
||||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||||
texts = text_splitter.split_documents(documents)
|
||||
|
||||
# Clean up the temporary file
|
||||
os.unlink(tmp_path)
|
||||
|
||||
return texts
|
||||
|
|
@ -143,26 +115,18 @@ COLLECTION_NAME = "cohere_rag"
|
|||
def create_vector_stores(texts):
|
||||
"""Create and populate vector store with documents."""
|
||||
try:
|
||||
# First, create the collection explicitly
|
||||
try:
|
||||
client.create_collection(
|
||||
collection_name=COLLECTION_NAME,
|
||||
vectors_config=VectorParams(
|
||||
size=1024, # Dimension for Cohere embed-english-v3.0
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
client.create_collection(collection_name=COLLECTION_NAME,
|
||||
vectors_config=VectorParams(size=1024,
|
||||
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
|
||||
|
||||
# Then initialize the vector store
|
||||
vector_store = QdrantVectorStore(
|
||||
client=client,
|
||||
collection_name=COLLECTION_NAME,
|
||||
embedding=embedding,
|
||||
)
|
||||
vector_store = QdrantVectorStore(client=client,
|
||||
collection_name=COLLECTION_NAME,
|
||||
embedding=embedding)
|
||||
|
||||
with st.spinner('Storing documents in Qdrant...'):
|
||||
vector_store.add_documents(texts)
|
||||
|
|
@ -174,91 +138,99 @@ def create_vector_stores(texts):
|
|||
st.error(f"Error in vector store creation: {str(e)}")
|
||||
return None
|
||||
|
||||
def create_fallback_agent():
|
||||
"""Create a LangGraph agent with DuckDuckGo search tool."""
|
||||
# Define the state schema using TypedDict
|
||||
class AgentState(TypedDict):
|
||||
"""State schema for the agent."""
|
||||
messages: List[HumanMessage | AIMessage | SystemMessage]
|
||||
is_last_step: bool
|
||||
|
||||
class RateLimitedDuckDuckGo(DuckDuckGoSearchRun):
|
||||
@retry(wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
stop=stop_after_attempt(3))
|
||||
def run(self, query: str) -> str:
|
||||
"""Run search with rate limiting."""
|
||||
try:
|
||||
sleep(2) # Add delay between requests
|
||||
return super().run(query)
|
||||
except Exception as e:
|
||||
if "Ratelimit" in str(e):
|
||||
sleep(5) # Longer delay on rate limit
|
||||
return super().run(query)
|
||||
raise e
|
||||
|
||||
def create_fallback_agent(chat_model: BaseLanguageModel):
|
||||
"""Create a LangGraph agent for web research."""
|
||||
|
||||
def web_research(query: str) -> str:
|
||||
"""Search the web for information about a query."""
|
||||
search = DuckDuckGoSearchRun()
|
||||
results = search.run(query)
|
||||
return f"Web search results: {results}"
|
||||
|
||||
"""Web search with result formatting."""
|
||||
try:
|
||||
search = DuckDuckGoSearchRun(num_results=5)
|
||||
results = search.run(query)
|
||||
return results
|
||||
except Exception as e:
|
||||
return f"Search failed: {str(e)}. Providing answer based on general knowledge."
|
||||
|
||||
tools = [web_research]
|
||||
|
||||
# Create agent with Cohere model
|
||||
agent = create_react_agent(
|
||||
chat_model, # Using the already initialized Cohere model
|
||||
tools=tools,
|
||||
)
|
||||
agent = create_react_agent(model=chat_model,
|
||||
tools=tools,
|
||||
debug=False)
|
||||
|
||||
return agent
|
||||
|
||||
def process_query(vectorstore, query) -> tuple[str, list]:
|
||||
"""Process a query using RAG with fallback to web search."""
|
||||
try:
|
||||
# First try vector store retrieval
|
||||
retriever = vectorstore.as_retriever(
|
||||
search_type="similarity_score_threshold",
|
||||
search_kwargs={
|
||||
"k": 10,
|
||||
"score_threshold": 0.7 # Only return relevant documents
|
||||
"score_threshold": 0.7
|
||||
}
|
||||
)
|
||||
|
||||
# Get relevant documents
|
||||
with st.spinner('Searching document database...'):
|
||||
relevant_docs = retriever.get_relevant_documents(query)
|
||||
relevant_docs = retriever.get_relevant_documents(query)
|
||||
|
||||
if relevant_docs:
|
||||
# Use RAG with document context
|
||||
retrieval_qa_prompt = hub.pull("langchain-ai/retrieval-qa-chat")
|
||||
combine_docs_chain = create_stuff_documents_chain(chat_model, retrieval_qa_prompt)
|
||||
retrieval_chain = create_retrieval_chain(retriever, combine_docs_chain)
|
||||
response = retrieval_chain.invoke({"input": query})
|
||||
return response['answer'], relevant_docs
|
||||
|
||||
combine_docs_chain = create_stuff_documents_chain(
|
||||
chat_model,
|
||||
retrieval_qa_prompt
|
||||
)
|
||||
|
||||
retrieval_chain = create_retrieval_chain(
|
||||
retriever,
|
||||
combine_docs_chain
|
||||
)
|
||||
|
||||
with st.spinner('Generating response from documents...'):
|
||||
response = retrieval_chain.invoke({"input": query})
|
||||
if not response or 'answer' not in response:
|
||||
raise ValueError("No response generated")
|
||||
|
||||
return response['answer'], relevant_docs
|
||||
else:
|
||||
# Fallback to web search using LangGraph agent
|
||||
st.info("No relevant documents found. Searching the web...")
|
||||
st.info("No relevant documents found. Searching web...")
|
||||
fallback_agent = create_fallback_agent(chat_model)
|
||||
|
||||
fallback_agent = create_fallback_agent()
|
||||
|
||||
with st.spinner('Searching web and generating response...'):
|
||||
# Prepare input for the agent
|
||||
with st.spinner('Researching...'):
|
||||
agent_input = {
|
||||
"messages": [
|
||||
("user", f"Please search and answer this question: {query}")
|
||||
]
|
||||
HumanMessage(content=f"""Please thoroughly research the question: '{query}' and provide a detailed and comprehensive response. Make sure to gather the latest information from credible sources. Minimum 400 words.""")
|
||||
],
|
||||
"is_last_step": False
|
||||
}
|
||||
|
||||
# Get agent response
|
||||
response = fallback_agent.invoke(agent_input)
|
||||
last_message = response["messages"][-1]
|
||||
config = {"recursion_limit": 100}
|
||||
|
||||
if isinstance(last_message, tuple):
|
||||
answer = last_message[1]
|
||||
else:
|
||||
answer = last_message.content
|
||||
|
||||
return f"Based on web search: {answer}", []
|
||||
try:
|
||||
response = fallback_agent.invoke(agent_input, config=config)
|
||||
|
||||
if isinstance(response, dict) and "messages" in response:
|
||||
last_message = response["messages"][-1]
|
||||
answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
|
||||
|
||||
return f"""Comprehensive Research Results:
|
||||
{answer}
|
||||
""", []
|
||||
|
||||
except Exception as agent_error:
|
||||
fallback_response = chat_model.invoke(f"Please provide a general answer to: {query}").content
|
||||
return f"Web search unavailable. General response: {fallback_response}", []
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"Error processing query: {str(e)}")
|
||||
return "I encountered an error processing your query. Please try again.", []
|
||||
st.error(f"Error: {str(e)}")
|
||||
return "I encountered an error. Please try rephrasing your question.", []
|
||||
|
||||
#post processing - strip, summarize along with formatted sources
|
||||
def post_process(answer, sources):
|
||||
"""Post-process the answer and format sources."""
|
||||
answer = answer.strip()
|
||||
|
|
@ -266,7 +238,7 @@ def post_process(answer, sources):
|
|||
# Summarize long answers
|
||||
if len(answer) > 500:
|
||||
summary_prompt = f"Summarize the following answer in 2-3 sentences: {answer}"
|
||||
summary = chat_model.invoke(summary_prompt).content # Changed from predict to invoke
|
||||
summary = chat_model.invoke(summary_prompt).content
|
||||
answer = f"{summary}\n\nFull Answer: {answer}"
|
||||
|
||||
formatted_sources = []
|
||||
|
|
@ -275,26 +247,25 @@ def post_process(answer, sources):
|
|||
formatted_sources.append(formatted_source)
|
||||
return answer, formatted_sources
|
||||
|
||||
st.title("RAG Agent with Cohere 🤖") # New heading
|
||||
st.title("RAG Agent with Cohere 🤖")
|
||||
|
||||
uploaded_file = st.file_uploader("Choose a PDF or Image File", type=["pdf", "jpg", "jpeg"])
|
||||
|
||||
if uploaded_file is not None:
|
||||
if uploaded_file is not None and 'processed_file' not in st.session_state:
|
||||
with st.spinner('Processing file... This may take a while for images.'):
|
||||
texts = process_document(uploaded_file)
|
||||
vectorstore = create_vector_stores(texts)
|
||||
if vectorstore:
|
||||
st.session_state.vectorstore = vectorstore
|
||||
st.session_state.processed_file = True
|
||||
st.success('File uploaded and processed successfully!')
|
||||
else:
|
||||
st.error('Failed to process file. Please try again.')
|
||||
|
||||
# Display chat history
|
||||
for message in st.session_state.chat_history:
|
||||
with st.chat_message(message["role"]):
|
||||
st.markdown(message["content"])
|
||||
|
||||
# Chat input
|
||||
if query := st.chat_input("Ask a question about the document:"):
|
||||
st.session_state.chat_history.append({"role": "user", "content": query})
|
||||
with st.chat_message("user"):
|
||||
|
|
@ -304,31 +275,24 @@ if query := st.chat_input("Ask a question about the document:"):
|
|||
with st.chat_message("assistant"):
|
||||
try:
|
||||
answer, sources = process_query(st.session_state.vectorstore, query)
|
||||
st.markdown(answer)
|
||||
|
||||
if sources: # Only post-process if we have sources
|
||||
processed_answer, formatted_sources = post_process(answer, sources)
|
||||
else:
|
||||
processed_answer, formatted_sources = answer, []
|
||||
|
||||
st.markdown(f"{processed_answer}")
|
||||
|
||||
if formatted_sources:
|
||||
if sources:
|
||||
with st.expander("Sources"):
|
||||
for source in formatted_sources:
|
||||
st.markdown(f"- {source}")
|
||||
for source in sources:
|
||||
st.markdown(f"- {source.page_content[:200]}...")
|
||||
|
||||
st.session_state.chat_history.append({
|
||||
"role": "assistant",
|
||||
"content": processed_answer
|
||||
"role": "assistant",
|
||||
"content": answer
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"Error: {str(e)}")
|
||||
st.info("Please try asking your question again.")
|
||||
else:
|
||||
st.error("Please upload a document first.")
|
||||
|
||||
|
||||
# Add to sidebar
|
||||
with st.sidebar:
|
||||
st.divider()
|
||||
col1, col2 = st.columns(2)
|
||||
|
|
@ -339,7 +303,6 @@ with st.sidebar:
|
|||
with col2:
|
||||
if st.button('Clear All Data'):
|
||||
try:
|
||||
# Check if collections exist before deleting
|
||||
collections = client.get_collections().collections
|
||||
collection_names = [col.name for col in collections]
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue