Update app.py

app.py

@@ -1,233 +1,340 @@
-# Drug Repurposing Advisor: A Multi-Agent Workflow Example
-# This example uses dummy data for demonstration.
-# In a production system, replace the dummy data with real pharmaceutical databases.
+# Advanced Multi‑Modal Agentic RAG Chatbot
+# pip install -r requirements.txt
 
-import os
-import json
-import requests
 import streamlit as st
-from typing import List, Union, Tuple
+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
-from typing_extensions import TypedDict, Annotated
-from typing import Sequence
 
-# Dummy data for drug mechanism research and clinical trial outcomes
-drug_mechanism_texts = [
-    "Drug A: Inhibits enzyme X and modulates receptor Y; potential anti-inflammatory effects.",
-    "Drug B: Blocks ion channel Z; has been shown to reduce oxidative stress in preclinical models.",
-    "Drug C: Activates nuclear receptor W; exhibits neuroprotective properties."
+# -------------------------------------------------------------------
+# 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."
 ]
 
-clinical_trials_texts = [
-    "Trial 1: Drug A repurposed for rheumatoid arthritis showed a 30% improvement in joint function.",
-    "Trial 2: Drug B evaluated in a pilot study for neurodegenerative disorders demonstrated a reduction in symptom severity.",
-    "Trial 3: Drug C tested in a phase II trial for multiple sclerosis reported significant reduction in relapse rates."
+# 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
+# Text splitting settings: adaptable for multi‑modal data (could extend to images)
 splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
-mechanism_docs = splitter.create_documents(drug_mechanism_texts)
-clinical_docs = splitter.create_documents(clinical_trials_texts)
-
-# Here you would typically create vector embeddings and vectorstores (e.g., using ChromaDB)
-# For demonstration, we define simple retriever functions that return dummy results.
-def mechanism_retriever(query: str) -> str:
-    # Dummy search: return first document that mentions a keyword from the query
-    for doc in drug_mechanism_texts:
-        if any(word.lower() in doc.lower() for word in query.split()):
-            return f"[Mechanism Doc]: {doc}"
-    return "No relevant mechanism data found."
-
-def clinical_retriever(query: str) -> str:
-    for doc in clinical_trials_texts:
-        if any(word.lower() in doc.lower() for word in query.split()):
-            return f"[Clinical Trial Doc]: {doc}"
-    return "No relevant clinical trial data found."
-
-# Define tools using a simple wrapper function
-def create_retriever_tool(retriever_func, tool_name: str, description: str):
-    def tool(query: str):
-        return retriever_func(query)
-    # Mimic a tool message (in a real system, you would wrap this in a ToolMessage object)
-    tool.__name__ = tool_name
-    tool.description = description
-    return tool
-
-mechanism_tool = create_retriever_tool(
-    mechanism_retriever,
-    "mechanism_db_tool",
-    "Search drug mechanism data for repurposing insights."
+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"
 )
-clinical_tool = create_retriever_tool(
-    clinical_retriever,
-    "clinical_db_tool",
-    "Search clinical trial outcomes for repurposing evidence."
+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"
 )
-tools = [mechanism_tool, clinical_tool]
 
-# Define the AgentState type for our workflow
-class AgentState(TypedDict):
-    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
+research_retriever = research_vectorstore.as_retriever()
+development_retriever = development_vectorstore.as_retriever()
 
-# Agent function: Classifies queries as targeting drug mechanisms or clinical outcomes
-def agent(state: AgentState):
-    print("---CALL AGENT---")
-    messages = state["messages"]
-    user_message = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]
+# 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]
 
-    # Build a prompt to classify the query
-    prompt = f"""Given the user question: "{user_message}"
-If the question is about the molecular mechanism or pharmacodynamics, respond EXACTLY in this format:
-SEARCH_MECHANISM: <search terms>
+# -------------------------------------------------------------------
+# AGENT DESIGN: Advanced Agent with Self‑Reflection & Multi‑Tool Coordination
 
-If it's about clinical trial outcomes, efficacy, or safety evidence, respond EXACTLY in this format:
-SEARCH_CLINICAL: <search terms>
+class AdvancedAgentState(TypedDict):
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
 
-Otherwise, answer directly with general repurposing insights.
+def advanced_agent(state: AdvancedAgentState):
     """
-    # (Here we simulate a call to DeepSeek-R1 using a dummy response)
-    # In a real implementation, make an API call to DeepSeek-R1.
-    if "mechanism" in user_message.lower() or "how it works" in user_message.lower():
-        response_text = f"SEARCH_MECHANISM: {user_message}"
-    elif "trial" in user_message.lower() or "efficacy" in user_message.lower() or "safety" in user_message.lower():
-        response_text = f"SEARCH_CLINICAL: {user_message}"
-    else:
-        response_text = "The system did not classify your query. Please rephrase to focus on drug mechanism or clinical data."
-    
-    print("Agent response:", response_text)
-    # Format response into expected tool call format
-    if "SEARCH_MECHANISM:" in response_text:
-        query = response_text.split("SEARCH_MECHANISM:")[1].strip()
-        result = mechanism_tool(query)
-        return {"messages": [AIMessage(content=f'Action: mechanism_db_tool\n{{"query": "{query}"}}\n\nResults: {result}')]}
-
-    elif "SEARCH_CLINICAL:" in response_text:
-        query = response_text.split("SEARCH_CLINICAL:")[1].strip()
-        result = clinical_tool(query)
-        return {"messages": [AIMessage(content=f'Action: clinical_db_tool\n{{"query": "{query}"}}\n\nResults: {result}')]}
+    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)]}
 
-# Grading function: Checks if retrieved documents were found
-def simple_grade_documents(state: AgentState):
+# -------------------------------------------------------------------
+# 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]
-    print("Evaluating message:", last_message.content)
-    if "Results:" in last_message.content and "No relevant" not in last_message.content:
-        print("---DATA FOUND, PROCEED TO GENERATE INSIGHTS---")
+    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:
-        print("---NO DATA FOUND, TRY REWRITE---")
+        st.write(">> [Grade] No sufficient documents; try rewriting the query.")
         return "rewrite"
 
-# Generate function: Synthesizes repurposing insights from retrieved data
-def generate(state: AgentState):
-    print("---GENERATE FINAL INSIGHTS---")
+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"]
-    question = messages[0].content
+    original_question = messages[0].content
     last_message = messages[-1]
-    # Extract data from results
-    data_start = last_message.content.find("Results:")
-    retrieved_data = last_message.content[data_start:] if data_start != -1 else "No data available"
-    # Build a prompt to synthesize insights
-    prompt = f"""Based on the following retrieved data:
-    {retrieved_data}
-and considering the question:
-    {question}
-Summarize potential drug repurposing opportunities and any recommended next steps for further investigation.
+    
+    # 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):
     """
-    # Dummy generation using a simple echo for demonstration.
-    final_answer = f"Summary Insight: Considering the data, a promising repurposing opportunity is to explore Drug A for anti-inflammatory applications beyond its original use, and Drug B might be repurposed for neurodegenerative conditions. Further research should validate these hypotheses."
-    print("Final Answer:", final_answer)
-    return {"messages": [AIMessage(content=final_answer)]}
-
-# Rewrite function: If no data is found, help rephrase the query for clarity
-def rewrite(state: AgentState):
-    print("---REWRITE QUESTION---")
+    Rewrite the user query for clarity using a self-reflection process.
+    """
+    st.write(">> [Rewrite] Improving query clarity...")
     messages = state["messages"]
-    original_question = messages[0].content if messages else "N/A"
-    # Dummy rewrite that just appends "Please specify mechanism or trial data" for demonstration.
-    rewritten = f"{original_question} (Please specify if you are asking about drug mechanism or clinical trial outcomes.)"
-    print("Rewritten question:", rewritten)
-    return {"messages": [AIMessage(content=rewritten)]}
-
-# Decision function: Determines next step based on last message content
-def custom_tools_condition(state: AgentState):
+    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
-    if content.startswith("Action:"):
-        print("Tool action detected. Proceed to retrieval.")
+    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
 
-# Create the workflow graph
-workflow = StateGraph(AgentState)
-workflow.add_node("agent", agent)
-retrieve_node = ToolNode(tools)
-workflow.add_node("retrieve", retrieve_node)
-workflow.add_node("rewrite", rewrite)
-workflow.add_node("generate", generate)
-
-# Define workflow edges
-workflow.add_edge(START, "agent")
-workflow.add_conditional_edges("agent", custom_tools_condition, {"tools": "retrieve", END: END})
-workflow.add_conditional_edges("retrieve", simple_grade_documents)
-workflow.add_edge("generate", END)
-workflow.add_edge("rewrite", "agent")
-app = workflow.compile()
-
-# Function to process a query through the workflow
-def process_question(user_question: str, app, config: dict):
+# -------------------------------------------------------------------
+# 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 for the Drug Repurposing Advisor
+# -------------------------------------------------------------------
+# STREAMLIT UI: Multi‑Modal Advanced Chatbot Interface
+
 def main():
     st.set_page_config(
-        page_title="Drug Repurposing Advisor",
+        page_title="Advanced Multi‑Modal AI Assistant",
         layout="wide",
         initial_sidebar_state="expanded"
     )
-    st.title("💊 Drug Repurposing Advisor")
-    st.markdown("### Explore potential drug repurposing opportunities with AI-driven insights.")
-    query = st.text_area("Enter your research question:",
-                          placeholder="e.g., Can Drug A be repurposed for neurodegenerative diseases?")
+    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 Insights", use_container_width=True):
+        if st.button("🔍 Get Advanced Answer", use_container_width=True):
             if query:
-                with st.spinner("Processing your query..."):
-                    events = process_question(query, app, {"configurable": {"thread_id": "1"}})
+                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 Processing Step", expanded=True):
-                                content = event['agent']['messages'][0].content
-                                st.markdown(f"**Agent Step Output:**\n\n{content}")
+                            with st.expander("🔄 Agent Analysis", expanded=True):
+                                st.info(event['agent']['messages'][0].content)
                         elif 'generate' in event:
-                            st.markdown("### Final Insights:")
+                            st.markdown("### ✨ Final Answer:")
                             st.success(event['generate']['messages'][0].content)
                         elif 'rewrite' in event:
-                            st.markdown("### Suggestion:")
-                            st.warning(event['rewrite']['messages'][0].content)
+                            st.warning("Query was unclear. Rewriting...")
+                            st.info(event['rewrite']['messages'][0].content)
             else:
-                st.warning("⚠️ Please enter a query.")
+                st.warning("⚠️ Please enter a question!")
     with col2:
         st.markdown("""
-        **How to Use the Drug Repurposing Advisor:**
-        1. **Input Query:** Describe your research question. Specify whether you are interested in drug mechanisms or clinical outcomes.
-        2. **Get Insights:** Click "Get Insights" and let the system process your query.
-        3. **Review Output:** Explore the retrieved data and the final synthesized insights.
-        **Example Questions:**
-        - "How does Drug A work and could its mechanism be useful in treating inflammatory diseases?"
-        - "What are the clinical trial outcomes of Drug B and can it be repurposed for neurodegenerative conditions?"
+        ### 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()