Update app.py

app.py

@@ -1,300 +1,233 @@
-# app.py
-# Advanced AI R&D Assistant for Hugging Face Spaces
-#
-# This app leverages LangGraph, DeepSeek-R1 via text-based function calling, and Agentic RAG.
-# API keys are securely loaded via environment variables.
-#
-# To deploy:
-# 1. Add your API key to Hugging Face Space secrets with the key DEEP_SEEK_API.
-# 2. Ensure your requirements.txt is properly configured.
-# 3. Run the app with Streamlit.
+# 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.
 
 import os
-import re
-import logging
-import streamlit as st
+import json
 import requests
-from typing import Sequence
-from typing_extensions import TypedDict, Annotated
-
-# Updated imports for LangChain
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores import Chroma
-from langchain.schema import HumanMessage, AIMessage, ToolMessage
+import streamlit as st
+from typing import List, Union, Tuple
+from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.tools.retriever import create_retriever_tool
-
-# Imports for LangGraph remain the same
 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
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# --- Dummy Data Setup ---
-research_texts = [
-    "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%",
-    "Academic Paper Summary: Why Transformers Became the Mainstream Architecture in Natural Language Processing",
-    "Latest Trends in Machine Learning Methods Using Quantum Computing"
+# 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."
 ]
 
-development_texts = [
-    "Project A: UI Design Completed, API Integration in Progress",
-    "Project B: Testing New Feature X, Bug Fixes Needed",
-    "Product Y: In the Performance Optimization Stage Before Release"
+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."
 ]
 
-# --- Preprocessing & Embeddings ---
+# Text splitting settings
 splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
-research_docs = splitter.create_documents(research_texts)
-development_docs = splitter.create_documents(development_texts)
-
-embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
-
-research_vectorstore = Chroma.from_documents(
-    documents=research_docs,
-    embedding=embeddings,
-    collection_name="research_collection"
-)
-development_vectorstore = Chroma.from_documents(
-    documents=development_docs,
-    embedding=embeddings,
-    collection_name="development_collection"
-)
-
-research_retriever = research_vectorstore.as_retriever()
-development_retriever = development_vectorstore.as_retriever()
-
-research_tool = create_retriever_tool(
-    research_retriever,
-    "research_db_tool",
-    "Search information from the research database."
+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."
 )
-development_tool = create_retriever_tool(
-    development_retriever,
-    "development_db_tool",
-    "Search information from the development database."
+clinical_tool = create_retriever_tool(
+    clinical_retriever,
+    "clinical_db_tool",
+    "Search clinical trial outcomes for repurposing evidence."
 )
-tools = [research_tool, development_tool]
+tools = [mechanism_tool, clinical_tool]
 
-# --- Agent and Workflow Functions ---
+# Define the AgentState type for our workflow
 class AgentState(TypedDict):
     messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
 
+# Agent function: Classifies queries as targeting drug mechanisms or clinical outcomes
 def agent(state: AgentState):
-    logger.info("Agent invoked")
+    print("---CALL AGENT---")
     messages = state["messages"]
-    user_message = messages[0][1] if isinstance(messages[0], tuple) else messages[0].content
-
-    prompt = f"""Given this user question: "{user_message}"
-If it's about research or academic topics, respond EXACTLY in this format:
-SEARCH_RESEARCH: <search terms>
-
-If it's about development status, respond EXACTLY in this format:
-SEARCH_DEV: <search terms>
+    user_message = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]
+
+    # 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>
+
+If it's about clinical trial outcomes, efficacy, or safety evidence, respond EXACTLY in this format:
+SEARCH_CLINICAL: <search terms>
+
+Otherwise, answer directly with general repurposing insights.
+    """
+    # (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}')]}
 
-Otherwise, just answer directly.
-"""
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {os.environ.get('DEEP_SEEK_API')}",
-        "Content-Type": "application/json"
-    }
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
-    response = requests.post(
-        "https://api.deepseek.com/v1/chat/completions",
-        headers=headers,
-        json=data,
-        verify=False
-    )
-    if response.status_code == 200:
-        response_text = response.json()['choices'][0]['message']['content']
-        logger.info(f"DeepSeek response: {response_text}")
-        if "SEARCH_RESEARCH:" in response_text:
-            query = response_text.split("SEARCH_RESEARCH:")[1].strip()
-            results = research_retriever.invoke(query)
-            return {"messages": [AIMessage(content=f'Action: research_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
-        elif "SEARCH_DEV:" in response_text:
-            query = response_text.split("SEARCH_DEV:")[1].strip()
-            results = development_retriever.invoke(query)
-            return {"messages": [AIMessage(content=f'Action: development_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
-        else:
-            return {"messages": [AIMessage(content=response_text)]}
     else:
-        error_msg = f"DeepSeek API call failed: {response.text}"
-        logger.error(error_msg)
-        raise Exception(error_msg)
+        return {"messages": [AIMessage(content=response_text)]}
 
+# Grading function: Checks if retrieved documents were found
 def simple_grade_documents(state: AgentState):
-    last_message = state["messages"][-1]
-    logger.info(f"Grading message: {last_message.content}")
-    if "Results: [Document" in last_message.content:
+    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---")
         return "generate"
     else:
+        print("---NO DATA FOUND, TRY REWRITE---")
         return "rewrite"
 
+# Generate function: Synthesizes repurposing insights from retrieved data
 def generate(state: AgentState):
-    logger.info("Generating final answer")
+    print("---GENERATE FINAL INSIGHTS---")
     messages = state["messages"]
-    question = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]
+    question = messages[0].content
     last_message = messages[-1]
-    docs = ""
-    if "Results: [" in last_message.content:
-        docs = last_message.content[last_message.content.find("Results: ["):]
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {os.environ.get('DEEP_SEEK_API')}",
-        "Content-Type": "application/json"
-    }
-    prompt = f"""Based on these research documents, summarize the latest advancements in AI:
-Question: {question}
-Documents: {docs}
-Focus on extracting and synthesizing the key findings from the research papers.
-"""
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
-    response = requests.post(
-        "https://api.deepseek.com/v1/chat/completions",
-        headers=headers,
-        json=data,
-        verify=False
-    )
-    if response.status_code == 200:
-        response_text = response.json()['choices'][0]['message']['content']
-        return {"messages": [AIMessage(content=response_text)]}
-    else:
-        error_msg = f"DeepSeek API generate call failed: {response.text}"
-        logger.error(error_msg)
-        raise Exception(error_msg)
-
+    # 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.
+    """
+    # 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):
-    logger.info("Rewriting question")
-    original_question = state["messages"][0].content if state["messages"] else "N/A"
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {os.environ.get('DEEP_SEEK_API')}",
-        "Content-Type": "application/json"
-    }
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": f"Rewrite this question to be more specific and clearer: {original_question}"}],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
-    response = requests.post(
-        "https://api.deepseek.com/v1/chat/completions",
-        headers=headers,
-        json=data,
-        verify=False
-    )
-    if response.status_code == 200:
-        response_text = response.json()['choices'][0]['message']['content']
-        return {"messages": [AIMessage(content=response_text)]}
-    else:
-        error_msg = f"DeepSeek API rewrite call failed: {response.text}"
-        logger.error(error_msg)
-        raise Exception(error_msg)
+    print("---REWRITE QUESTION---")
+    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)]}
 
-tools_pattern = re.compile(r"Action: .*")
+# Decision function: Determines next step based on last message content
 def custom_tools_condition(state: AgentState):
-    last_message = state["messages"][-1]
-    if tools_pattern.match(last_message.content):
+    messages = state["messages"]
+    last_message = messages[-1]
+    content = last_message.content
+    if content.startswith("Action:"):
+        print("Tool action detected. Proceed to retrieval.")
         return "tools"
     return END
 
-# Build the workflow with LangGraph's StateGraph
+# 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 = workflow.compile()
+app = workflow.compile()
 
-def process_question(user_question, app, config):
+# Function to process a query through the workflow
+def process_question(user_question: str, app, config: dict):
     events = []
     for event in app.stream({"messages": [("user", user_question)]}, config):
         events.append(event)
     return events
 
-# --- Streamlit UI ---
+# Streamlit UI for the Drug Repurposing Advisor
 def main():
-    st.set_page_config(page_title="Advanced AI R&D Assistant", layout="wide", initial_sidebar_state="expanded")
-    st.markdown(
-        """
-        <style>
-        .stApp { background-color: #f8f9fa; }
-        .stButton > button { width: 100%; margin-top: 20px; }
-        .data-box { padding: 20px; border-radius: 10px; margin: 10px 0; }
-        .research-box { background-color: #e3f2fd; border-left: 5px solid #1976d2; }
-        .dev-box { background-color: #e8f5e9; border-left: 5px solid #43a047; }
-        </style>
-        """, unsafe_allow_html=True
+    st.set_page_config(
+        page_title="Drug Repurposing Advisor",
+        layout="wide",
+        initial_sidebar_state="expanded"
     )
-
-    # Sidebar: Display available data
-    with st.sidebar:
-        st.header("📚 Available Data")
-        st.subheader("Research Database")
-        for text in research_texts:
-            st.markdown(f'<div class="data-box research-box">{text}</div>', unsafe_allow_html=True)
-        st.subheader("Development Database")
-        for text in development_texts:
-            st.markdown(f'<div class="data-box dev-box">{text}</div>', unsafe_allow_html=True)
-
-    st.title("🤖 Advanced AI R&D Assistant")
-    st.markdown("---")
-    query = st.text_area("Enter your question:", height=100, placeholder="e.g., What is the latest advancement in AI research?")
-
+    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?")
     col1, col2 = st.columns([1, 2])
     with col1:
-        if st.button("🔍 Get Answer", use_container_width=True):
+        if st.button("🔍 Get Insights", use_container_width=True):
             if query:
-                with st.spinner('Processing your question...'):
-                    events = process_question(query, app_workflow, {"configurable": {"thread_id": "1"}})
+                with st.spinner("Processing your query..."):
+                    events = process_question(query, app, {"configurable": {"thread_id": "1"}})
                     for event in events:
                         if 'agent' in event:
-                            with st.expander("🔄 Processing Step", expanded=True):
+                            with st.expander("Agent Processing Step", expanded=True):
                                 content = event['agent']['messages'][0].content
-                                if "Results:" in content:
-                                    st.markdown("### 📑 Retrieved Documents:")
-                                    docs = content[content.find("Results:"):]
-                                    st.info(docs)
+                                st.markdown(f"**Agent Step Output:**\n\n{content}")
                         elif 'generate' in event:
-                            st.markdown("### ✨ Final Answer:")
+                            st.markdown("### Final Insights:")
                             st.success(event['generate']['messages'][0].content)
+                        elif 'rewrite' in event:
+                            st.markdown("### Suggestion:")
+                            st.warning(event['rewrite']['messages'][0].content)
             else:
-                st.warning("⚠️ Please enter a question first!")
+                st.warning("⚠️ Please enter a query.")
     with col2:
-        st.markdown(
-            """
-            ### 🎯 How to Use
-            1. Type your question in the text box.
-            2. Click "Get Answer" to process.
-            3. View retrieved documents and the final answer.
-            
-            ### 💡 Example Questions
-            - What are the latest advancements in AI research?
-            - What is the status of Project A?
-            - What are the current trends in machine learning?
-            """
-        )
-
+        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?"
+        """)
 if __name__ == "__main__":
     main()