app.py

# Advanced Multi‑Modal Agentic RAG Chatbot
# pip install -r requirements.txt

import streamlit as st
import requests
import json
import re
import os
from typing import Sequence
from typing_extensions import TypedDict, Annotated

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langgraph.graph import END, StateGraph, START
from langgraph.prebuilt import ToolNode
from langgraph.graph.message import add_messages

# -------------------------------------------------------------------
# DATA SETUP: Static (research) and Dynamic (live updates) Databases

# Static research data (e.g., academic papers, reports)
research_texts = [
    "Research Report: New algorithm boosts image recognition to 99%.",
    "Paper: Transformers have redefined natural language processing paradigms.",
    "Deep dive: Quantum computing’s emerging role in machine learning."
]

# Dynamic development/live data (e.g., real-time project updates)
development_texts = [
    "Live Update: Project X API integration at 75% completion.",
    "Status: Project Y is undergoing stress testing for scalability.",
    "Alert: Immediate patch required for Project Z deployment issues."
]

# Text splitting settings: adaptable for multi‑modal data (could extend to images)
splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
research_docs = splitter.create_documents(research_texts)
development_docs = splitter.create_documents(development_texts)

# Create vector stores using advanced embeddings
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-large"
)
research_vectorstore = Chroma.from_documents(
    documents=research_docs,
    embedding=embeddings,
    collection_name="research_collection_adv"
)
development_vectorstore = Chroma.from_documents(
    documents=development_docs,
    embedding=embeddings,
    collection_name="development_collection_adv"
)

research_retriever = research_vectorstore.as_retriever()
development_retriever = development_vectorstore.as_retriever()

# Create tool wrappers for the two databases
from langchain.tools.retriever import create_retriever_tool
research_tool = create_retriever_tool(
    research_retriever,
    "research_db_tool",
    "Search and retrieve static academic research documents."
)
development_tool = create_retriever_tool(
    development_retriever,
    "development_db_tool",
    "Retrieve dynamic, real‑time development updates."
)
tools = [research_tool, development_tool]

# -------------------------------------------------------------------
# AGENT DESIGN: Advanced Agent with Self‑Reflection & Multi‑Tool Coordination

class AdvancedAgentState(TypedDict):
    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]

def advanced_agent(state: AdvancedAgentState):
    """
    A smarter agent that:
    • Receives a multi-modal query (text and potentially images)
    • Self-reflects on the query to decide if a real-time lookup is needed
    • Chooses the appropriate tool or even combines results if required.
    """
    st.write(">> [Agent] Processing query...")
    messages = state["messages"]
    user_message = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]

    # Step 1: Initial Analysis and Self-Reflection
    analysis_prompt = f"""You are an advanced multi-modal reasoning engine.
User Query: "{user_message}"
Analyze the query and decide:
- If it is about static academic research, output EXACTLY: ACTION_RESEARCH: <query>.
- If it is about dynamic development or live updates, output EXACTLY: ACTION_LIVE: <query>.
- Otherwise, output a direct answer with self-reflection.
Also, add a brief self-reflection on your reasoning process.
"""
    headers = {
        "Accept": "application/json",
        "Authorization": "Bearer sk-ADVANCEDKEY123",  # Use your secure key here
        "Content-Type": "application/json"
    }
    data = {
        "model": "deepseek-chat",
        "messages": [{"role": "user", "content": analysis_prompt}],
        "temperature": 0.6,
        "max_tokens": 1024
    }
    response = requests.post(
        "https://api.deepseek.com/v1/chat/completions",
        headers=headers,
        json=data,
        verify=False
    )
    if response.status_code != 200:
        raise Exception(f"API call failed: {response.text}")
    response_text = response.json()['choices'][0]['message']['content']
    st.write(">> [Agent] Analysis:", response_text)
    
    # Step 2: Interpret the result and call the appropriate tool(s)
    if "ACTION_RESEARCH:" in response_text:
        query = response_text.split("ACTION_RESEARCH:")[1].strip().split("\n")[0]
        results = research_retriever.invoke(query)
        return {"messages": [AIMessage(content=f'Action: research_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}\n\nReflection: {response_text}')]}
    elif "ACTION_LIVE:" in response_text:
        query = response_text.split("ACTION_LIVE:")[1].strip().split("\n")[0]
        results = development_retriever.invoke(query)
        return {"messages": [AIMessage(content=f'Action: development_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}\n\nReflection: {response_text}')]}
    else:
        # Direct answer with self-reflection
        return {"messages": [AIMessage(content=response_text)]}

# -------------------------------------------------------------------
# DECISION & GENERATION FUNCTIONS: Advanced Grading & Iterative Answering

def advanced_grade(state: AdvancedAgentState):
    """
    Checks the last message for valid document retrieval or if further refinement is needed.
    """
    messages = state["messages"]
    last_message = messages[-1]
    st.write(">> [Grade] Reviewing output:", last_message.content)
    if "Results: [Document" in last_message.content:
        st.write(">> [Grade] Documents found; proceed to generation.")
        return "generate"
    else:
        st.write(">> [Grade] No sufficient documents; try rewriting the query.")
        return "rewrite"

def advanced_generate(state: AdvancedAgentState):
    """
    Generate a final answer by summarizing retrieved documents
    while incorporating self-reflection from the agent.
    """
    st.write(">> [Generate] Synthesizing final answer...")
    messages = state["messages"]
    original_question = messages[0].content
    last_message = messages[-1]
    
    # Extract retrieved documents if available
    docs = ""
    if "Results: [" in last_message.content:
        docs = last_message.content[last_message.content.find("Results: ["):]
    
    generate_prompt = f"""Using the following documents and the query below,
summarize a comprehensive answer.
Query: {original_question}
Documents: {docs}
Additionally, integrate the self-reflection notes from the agent to explain your reasoning.
Focus on clarity and depth.
"""
    headers = {
        "Accept": "application/json",
        "Authorization": "Bearer sk-ADVANCEDKEY123",
        "Content-Type": "application/json"
    }
    data = {
        "model": "deepseek-chat",
        "messages": [{"role": "user", "content": generate_prompt}],
        "temperature": 0.65,
        "max_tokens": 1024
    }
    response = requests.post(
        "https://api.deepseek.com/v1/chat/completions",
        headers=headers,
        json=data,
        verify=False
    )
    if response.status_code != 200:
        raise Exception(f"API call failed during generation: {response.text}")
    final_text = response.json()['choices'][0]['message']['content']
    st.write(">> [Generate] Final Answer generated.")
    return {"messages": [AIMessage(content=final_text)]}

def advanced_rewrite(state: AdvancedAgentState):
    """
    Rewrite the user query for clarity using a self-reflection process.
    """
    st.write(">> [Rewrite] Improving query clarity...")
    messages = state["messages"]
    original_query = messages[0].content
    headers = {
        "Accept": "application/json",
        "Authorization": "Bearer sk-ADVANCEDKEY123",
        "Content-Type": "application/json"
    }
    data = {
        "model": "deepseek-chat",
        "messages": [{"role": "user", "content": f"Please rewrite this query for more specificity and clarity: {original_query}"}],
        "temperature": 0.6,
        "max_tokens": 1024
    }
    response = requests.post(
        "https://api.deepseek.com/v1/chat/completions",
        headers=headers,
        json=data,
        verify=False
    )
    if response.status_code != 200:
        raise Exception(f"API call failed during rewrite: {response.text}")
    rewritten_query = response.json()['choices'][0]['message']['content']
    st.write(">> [Rewrite] Rewritten query:", rewritten_query)
    return {"messages": [AIMessage(content=rewritten_query)]}

# -------------------------------------------------------------------
# Custom Tools Condition: Advanced Multi‑Tool Routing

advanced_tools_pattern = re.compile(r"Action: .*")

def advanced_tools_condition(state: AdvancedAgentState):
    messages = state["messages"]
    last_message = messages[-1]
    content = last_message.content
    st.write(">> [Condition] Checking for tool invocation:", content)
    if advanced_tools_pattern.match(content):
        st.write(">> [Condition] Routing to tools retrieval.")
        return "tools"
    st.write(">> [Condition] No tool call detected; ending workflow.")
    return END

# -------------------------------------------------------------------
# BUILDING THE ADVANCED WORKFLOW WITH LANGGRAPH

advanced_workflow = StateGraph(AdvancedAgentState)
advanced_workflow.add_node("agent", advanced_agent)
advanced_tool_node = ToolNode(tools)  # Re-use our existing tools
advanced_workflow.add_node("retrieve", advanced_tool_node)
advanced_workflow.add_node("rewrite", advanced_rewrite)
advanced_workflow.add_node("generate", advanced_generate)

advanced_workflow.add_edge(START, "agent")
advanced_workflow.add_conditional_edges(
    "agent",
    advanced_tools_condition,
    {"tools": "retrieve", END: END}
)
advanced_workflow.add_conditional_edges("retrieve", advanced_grade)
advanced_workflow.add_edge("generate", END)
advanced_workflow.add_edge("rewrite", "agent")

advanced_app = advanced_workflow.compile()

def process_advanced_question(user_question, app, config):
    """Process user question through the advanced workflow."""
    events = []
    for event in app.stream({"messages": [("user", user_question)]}, config):
        events.append(event)
    return events

# -------------------------------------------------------------------
# STREAMLIT UI: Multi‑Modal Advanced Chatbot Interface

def main():
    st.set_page_config(
        page_title="Advanced Multi‑Modal AI Assistant",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    st.markdown("""
    <style>
    .stApp { background-color: #f0f2f6; }
    .stButton > button { width: 100%; margin-top: 20px; }
    .data-box { padding: 15px; border-radius: 8px; margin: 8px 0; }
    .research-box { background-color: #e1f5fe; border-left: 5px solid #0288d1; }
    .live-box { background-color: #e8f5e9; border-left: 5px solid #2e7d32; }
    </style>
    """, unsafe_allow_html=True)

    # Sidebar: Display static and live data
    with st.sidebar:
        st.header("📚 Data Sources")
        st.subheader("Static Research")
        for text in research_texts:
            st.markdown(f'<div class="data-box research-box">{text}</div>', unsafe_allow_html=True)
        st.subheader("Live Updates")
        for text in development_texts:
            st.markdown(f'<div class="data-box live-box">{text}</div>', unsafe_allow_html=True)

    st.title("🤖 Advanced Multi‑Modal Agentic RAG Assistant")
    st.markdown("---")
    
    # Query Input (supports future multi‑modal extensions)
    query = st.text_area("Enter your question (or upload an image in future versions):", height=100, placeholder="e.g., What recent breakthroughs in AI are influencing real‑time projects?")
    
    col1, col2 = st.columns([1, 2])
    with col1:
        if st.button("🔍 Get Advanced Answer", use_container_width=True):
            if query:
                with st.spinner("Processing your advanced query..."):
                    events = process_advanced_question(query, advanced_app, {"configurable": {"thread_id": "advanced1"}})
                    for event in events:
                        if 'agent' in event:
                            with st.expander("🔄 Agent Analysis", expanded=True):
                                st.info(event['agent']['messages'][0].content)
                        elif 'generate' in event:
                            st.markdown("### ✨ Final Answer:")
                            st.success(event['generate']['messages'][0].content)
                        elif 'rewrite' in event:
                            st.warning("Query was unclear. Rewriting...")
                            st.info(event['rewrite']['messages'][0].content)
            else:
                st.warning("⚠️ Please enter a question!")
    with col2:
        st.markdown("""
        ### How It Works:
        1. **Advanced Agent**: Uses self-reflection to decide between static or live data.
        2. **Tool Coordination**: Routes queries to the appropriate retrieval tool.
        3. **Self‑Reflection & Iteration**: If retrieval fails, the query is rewritten for clarity.
        4. **Final Synthesis**: Retrieved documents are summarized into a final, clear answer.
        
        ### Example Queries:
        - "What new breakthroughs in quantum machine learning are there?"
        - "Provide live updates on the progress of Project X."
        - "Summarize the recent advancements in transformer models."
        """)

if __name__ == "__main__":
    main()