Upload app.py

app.py

@@ -0,0 +1,552 @@
+import os
+import streamlit as st
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from langchain_core.messages import BaseMessage, HumanMessage
+from typing import Annotated, List, Optional, Dict
+from typing_extensions import TypedDict
+from langchain_community.document_loaders import WebBaseLoader
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_core.tools import tool
+from langchain.agents import AgentExecutor, create_openai_functions_agent
+from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, StateGraph, START
+import functools
+import operator
+import logging
+import time
+from tenacity import retry, stop_after_attempt, wait_exponential, RetryError
+from pydantic import ValidationError
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize temporary directory
+if 'working_directory' not in st.session_state:
+    _TEMP_DIRECTORY = TemporaryDirectory()
+    st.session_state.working_directory = Path(_TEMP_DIRECTORY.name)
+
+WORKING_DIRECTORY = st.session_state.working_directory
+
+# Streamlit UI
+st.set_page_config(page_title="MARS: Multi-Agent Report Synthesizer", layout="wide")
+
+# Custom CSS for styling
+st.markdown("""
+    <style>
+        body {
+            background-color: #f5f5f5;
+            color: #333333;
+            font-family: 'Comic Sans MS', 'Comic Sans', cursive;
+        }
+        .report-container {
+            border-radius: 10px;
+            background-color: #ffcccb;
+            padding: 20px;
+        }
+        .sidebar .sidebar-content {
+            background-color: #333333;
+            color: #ffffff;
+        }
+        .stButton button {
+            background-color: #ff6347;
+            color: #ffffff;
+            border-radius: 5px;
+            font-size: 18px;
+            padding: 10px 20px;
+            font-weight: bold;
+        }
+        .stTextInput input {
+            border-radius: 5px;
+            border: 2px solid #ff6347;
+            font-size: 16px;
+            padding: 10px;
+            width: 100%;
+        }
+        .stTextInput label {
+            font-size: 18px;
+            font-weight: bold;
+            color: #333333;
+        }
+        .stSelectbox label, .stDownloadButton label {
+            font-size: 18px;
+            font-weight: bold;
+            color: #333333;
+        }
+        .stSelectbox div, .stDownloadButton div {
+            background-color: #ffcccb;
+            color: #333333;
+            border-radius: 5px;
+            padding: 10px;
+            font-size: 16px;
+        }
+    </style>
+""", unsafe_allow_html=True)
+
+st.title("🚀 MARS: Multi-agent Report Synthesizer 🤖")
+st.sidebar.title("📋 Instructions")
+st.sidebar.write("""
+1. Enter your query in the input box.
+2. Marvin AI will assign tasks to different teams.
+3. You can see the progress and download the final report.
+4. Use the buttons to list and download output files.
+""")
+
+# Input fields for API keys
+openai_api_key = st.sidebar.text_input("OpenAI API Key", type="password")
+tavily_api_key = st.sidebar.text_input("Tavily API Key", type="password")
+
+# Store the API keys in the session state
+if openai_api_key:
+    os.environ["OPENAI_API_KEY"] = openai_api_key
+if tavily_api_key:
+    os.environ["TAVILY_API_KEY"] = tavily_api_key
+
+# Check if the API keys are set
+if not os.getenv("OPENAI_API_KEY"):
+    st.error("OpenAI API Key is required.")
+if not os.getenv("TAVILY_API_KEY"):
+    st.error("Tavily API Key is required.")
+
+# Define tools
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
+def tavily_search_with_retry(*args, **kwargs):
+    try:
+        result = TavilySearchResults(*args, **kwargs)
+        return result
+    except ValidationError as ve:
+        logger.error(f"Validation error: {ve}")
+        raise ve
+    except Exception as e:
+        logger.error(f"Error in Tavily search: {e}")
+        raise e
+
+tavily_tool = tavily_search_with_retry(max_results=5)
+
+@tool
+def scrape_webpages(urls: List[str]) -> str:
+    """Use requests and bs4 to scrape the provided web pages for detailed information."""
+    try:
+        loader = WebBaseLoader(urls)
+        docs = loader.load()
+        return "\n\n".join(
+            [
+                f'\n{doc.page_content}\n'
+                for doc in docs
+            ]
+        )
+    except Exception as e:
+        logger.error(f"Error in scrape_webpages: {str(e)}")
+        return f"Error occurred while scraping webpages: {str(e)}"
+
+@tool
+def create_outline(
+    points: Annotated[List[str], "List of main points or sections."],
+    file_name: Annotated[str, "File path to save the outline."],
+) -> Annotated[str, "Path of the saved outline file."]:
+    """Create and save an outline."""
+    try:
+        with (WORKING_DIRECTORY / file_name).open("w") as file:
+            for i, point in enumerate(points):
+                file.write(f"{i + 1}. {point}\n")
+        return f"Outline saved to {file_name}"
+    except Exception as e:
+        logger.error(f"Error in create_outline: {str(e)}")
+        return f"Error occurred while creating outline: {str(e)}"
+
+@tool
+def read_document(
+    file_name: Annotated[str, "File path to save the document."],
+    start: Annotated[Optional[int], "The start line. Default is 0"] = None,
+    end: Annotated[Optional[int], "The end line. Default is None"] = None,
+) -> str:
+    """Read the specified document."""
+    try:
+        with (WORKING_DIRECTORY / file_name).open("r") as file:
+            lines = file.readlines()
+        if start is not None:
+            start = 0
+        return "\n".join(lines[start:end])
+    except Exception as e:
+        logger.error(f"Error in read_document: {str(e)}")
+        return f"Error occurred while reading document: {str(e)}"
+
+@tool
+def write_document(
+    content: Annotated[str, "Text content to be written into the document."],
+    file_name: Annotated[str, "File path to save the document."],
+) -> Annotated[str, "Path of the saved document file."]:
+    """Create and save a text document."""
+    try:
+        with (WORKING_DIRECTORY / file_name).open("w") as file:
+            file.write(content)
+        return f"Document saved to {file_name}"
+    except Exception as e:
+        logger.error(f"Error in write_document: {str(e)}")
+        return f"Error occurred while writing document: {str(e)}"
+
+@tool
+def edit_document(
+    file_name: Annotated[str, "Path of the document to be edited."],
+    inserts: Annotated[
+        Dict[int, str],
+        "Dictionary where key is the line number (1-indexed) and value is the text to be inserted at that line.",
+    ],
+) -> Annotated[str, "Path of the edited document file."]:
+    """Edit a document by inserting text at specific line numbers."""
+    try:
+        with (WORKING_DIRECTORY / file_name).open("r") as file:
+            lines = file.readlines()
+        sorted_inserts = sorted(inserts.items())
+        for line_number, text in sorted_inserts:
+            if 1 <= line_number <= len(lines) + 1:
+                lines.insert(line_number - 1, text + "\n")
+            else:
+                return f"Error: Line number {line_number} is out of range."
+        with (WORKING_DIRECTORY / file_name).open("w") as file:
+            file.writelines(lines)
+        return f"Document edited and saved to {file_name}"
+    except Exception as e:
+        logger.error(f"Error in edit_document: {str(e)}")
+        return f"Error occurred while editing document: {str(e)}"
+
+# Define the agents and their tools
+llm = ChatOpenAI(model="gpt-3.5-turbo-0125")
+
+def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str) -> str:
+    """Create a function-calling agent and add it to the graph."""
+    system_prompt += """\nWork autonomously according to your specialty, using the tools available to you.
+    Do not ask for clarification.
+    Your other team members (and other teams) will collaborate with you with their own specialties.
+    You are chosen for a reason! You are one of the following team members: {team_members}."""
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+    )
+    agent = create_openai_functions_agent(llm, tools, prompt)
+    executor = AgentExecutor(agent=agent, tools=tools)
+    return executor
+
+def agent_node(state, agent, name):
+    try:
+        logger.info(f"Starting {name} agent")
+        result = agent.invoke(state)
+        logger.info(f"{name} agent completed with result: {result}")
+        return {"messages": [HumanMessage(content=result["output"], name=name)]}
+    except ValidationError as ve:
+        logger.error(f"Validation error in {name} agent: {ve}")
+        return {"messages": [HumanMessage(content=f"Validation error in {name} agent: {ve}", name=name)]}
+    except Exception as e:
+        logger.error(f"Error in {name} agent: {e}")
+        return {"messages": [HumanMessage(content=f"Error occurred in {name} agent: {e}", name=name)]}
+
+def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> str:
+    """An LLM-based router."""
+    options = ["FINISH"] + members
+    function_def = {
+        "name": "route",
+        "description": "Select the next role.",
+        "parameters": {
+            "title": "routeSchema",
+            "type": "object",
+            "properties": {
+                "next": {
+                    "title": "Next",
+                    "anyOf": [
+                        {"enum": options},
+                    ],
+                },
+            },
+            "required": ["next"],
+        },
+    }
+    system_prompt += "\nEnsure that you direct the workflow to completion. If no progress is being made, or if the task seems complete, choose FINISH."
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder(variable_name="messages"),
+            ("system", "Given the conversation above, who should act next? Or should we FINISH? Select one of: {options}"),
+        ]
+    ).partial(options=str(options), team_members=", ".join(members))
+    return (
+        prompt
+        | llm.bind_functions(functions=[function_def], function_call="route")
+        | JsonOutputFunctionsParser()
+    )
+
+# ResearchTeam graph state
+class ResearchTeamState(TypedDict):
+    messages: Annotated[List[BaseMessage], operator.add]
+    team_members: List[str]
+    next: str
+
+llm = ChatOpenAI(model="gpt-3.5-turbo-0125")
+
+search_agent = create_agent(
+    llm,
+    [tavily_tool],
+    "You are a research assistant who can search for up-to-date info using the tavily search engine.",
+)
+search_node = functools.partial(agent_node, agent=search_agent, name="Search")
+
+research_agent = create_agent(
+    llm,
+    [scrape_webpages],
+    "You are a research assistant who can scrape specified urls for more detailed information using the scrape_webpages function.",
+)
+research_node = functools.partial(agent_node, agent=research_agent, name="WebScraper")
+
+supervisor_agent = create_team_supervisor(
+    llm,
+    "You are a supervisor tasked with managing a conversation between the"
+    " following workers: Search, WebScraper. Given the following user request,"
+    " respond with the worker to act next. Each worker will perform a"
+    " task and respond with their results and status. When finished,"
+    " respond with FINISH.",
+    ["Search", "WebScraper"],
+)
+
+research_graph = StateGraph(ResearchTeamState)
+research_graph.add_node("Search", search_node)
+research_graph.add_node("WebScraper", research_node)
+research_graph.add_node("supervisor", supervisor_agent)
+
+# Define the control flow
+research_graph.add_edge("Search", "supervisor")
+research_graph.add_edge("WebScraper", "supervisor")
+research_graph.add_conditional_edges(
+    "supervisor",
+    lambda x: x["next"],
+    {"Search": "Search", "WebScraper": "WebScraper", "FINISH": END},
+)
+
+research_graph.add_edge(START, "supervisor")
+chain = research_graph.compile()
+
+def enter_chain(message: str):
+    results = {
+        "messages": [HumanMessage(content=message)],
+    }
+    return results
+
+research_chain = enter_chain | chain
+
+# Document writing team graph state
+class DocWritingState(TypedDict):
+    messages: Annotated[List[BaseMessage], operator.add]
+    team_members: str
+    next: str
+    current_files: str
+
+def prelude(state):
+    written_files = []
+    if not WORKING_DIRECTORY.exists():
+        WORKING_DIRECTORY.mkdir()
+    try:
+        written_files = [
+            f.relative_to(WORKING_DIRECTORY) for f in WORKING_DIRECTORY.rglob("*")
+        ]
+    except Exception:
+        pass
+    if not written_files:
+        return {**state, "current_files": "No files written."}
+    return {
+        **state,
+        "current_files": "\nBelow are files your team has written to the directory:\n"
+        + "\n".join([f" - {f}" for f in written_files]),
+    }
+
+doc_writer_agent = create_agent(
+    llm,
+    [write_document, edit_document, read_document],
+    "You are an expert writing a research document.\n"
+    "Below are files currently in your directory:\n{current_files}",
+)
+context_aware_doc_writer_agent = prelude | doc_writer_agent
+doc_writing_node = functools.partial(
+    agent_node, agent=context_aware_doc_writer_agent, name="DocWriter"
+)
+
+note_taking_agent = create_agent(
+    llm,
+    [create_outline, read_document],
+    "You are an expert senior researcher tasked with writing a paper outline and"
+    " taking notes to craft a perfect paper.{current_files}",
+)
+context_aware_note_taking_agent = prelude | note_taking_agent
+note_taking_node = functools.partial(
+    agent_node, agent=context_aware_note_taking_agent, name="NoteTaker"
+)
+
+chart_generating_agent = create_agent(
+    llm,
+    [read_document],
+    "You are a data viz expert tasked with generating charts for a research project."
+    "{current_files}",
+)
+context_aware_chart_generating_agent = prelude | chart_generating_agent
+chart_generating_node = functools.partial(
+    agent_node, agent=context_aware_note_taking_agent, name="ChartGenerator"
+)
+
+doc_writing_supervisor = create_team_supervisor(
+    llm,
+    "You are a supervisor tasked with managing a conversation between the"
+    " following workers: {team_members}. Given the following user request,"
+    " respond with the worker to act next. Each worker will perform a"
+    " task and respond with their results and status. When finished,"
+    " respond with FINISH.",
+    ["DocWriter", "NoteTaker", "ChartGenerator"],
+)
+
+authoring_graph = StateGraph(DocWritingState)
+authoring_graph.add_node("DocWriter", doc_writing_node)
+authoring_graph.add_node("NoteTaker", note_taking_node)
+authoring_graph.add_node("ChartGenerator", chart_generating_node)
+authoring_graph.add_node("supervisor", doc_writing_supervisor)
+
+authoring_graph.add_edge("DocWriter", "supervisor")
+authoring_graph.add_edge("NoteTaker", "supervisor")
+authoring_graph.add_edge("ChartGenerator", "supervisor")
+authoring_graph.add_conditional_edges(
+    "supervisor",
+    lambda x: x["next"],
+    {
+        "DocWriter": "DocWriter",
+        "NoteTaker": "NoteTaker",
+        "ChartGenerator": "ChartGenerator",
+        "FINISH": END,
+    },
+)
+
+authoring_graph.add_edge(START, "supervisor")
+chain = authoring_graph.compile()
+
+def enter_chain(message: str, members: List[str]):
+    results = {
+        "messages": [HumanMessage(content=message)],
+        "team_members": ", ".join(members),
+    }
+    return results
+
+authoring_chain = (
+    functools.partial(enter_chain, members=authoring_graph.nodes)
+    | authoring_graph.compile()
+)
+
+llm = ChatOpenAI(model="gpt-3.5-turbo-0125")
+
+supervisor_node = create_team_supervisor(
+    llm,
+    "You are a supervisor tasked with managing a conversation between the"
+    " following teams: {team_members}. Given the following user request,"
+    " respond with the worker to act next. Each worker will perform a"
+    " task and respond with their results and status. Make sure each team is used atleast once. When finished,"
+    " respond with FINISH.",
+    ["ResearchTeam", "PaperWritingTeam"],
+)
+
+class State(TypedDict):
+    messages: Annotated[List[BaseMessage], operator.add]
+    next: str
+
+def get_last_message(state: State) -> str:
+    return state["messages"][-1].content
+
+def join_graph(response: dict):
+    return {"messages": [response["messages"][-1]]}
+
+super_graph = StateGraph(State)
+super_graph.add_node("ResearchTeam", get_last_message | research_chain | join_graph)
+super_graph.add_node("PaperWritingTeam", get_last_message | authoring_chain | join_graph)
+super_graph.add_node("supervisor", supervisor_node)
+
+super_graph.add_edge("ResearchTeam", "supervisor")
+super_graph.add_edge("PaperWritingTeam", "supervisor")
+super_graph.add_conditional_edges(
+    "supervisor",
+    lambda x: x["next"],
+    {
+        "PaperWritingTeam": "PaperWritingTeam",
+        "ResearchTeam": "ResearchTeam",
+        "FINISH": END,
+    },
+)
+super_graph.add_edge(START, "supervisor")
+super_graph = super_graph.compile()
+
+input_text = st.text_input("Enter your query:")
+
+if input_text and os.getenv("OPENAI_API_KEY") and os.getenv("TAVILY_API_KEY"):
+    st.markdown("### 🛠️ Task Progress")
+    start_time = time.time()
+    max_execution_time = 300  # 5 minutes
+
+    try:
+        for s in super_graph.stream(
+            {
+                "messages": [
+                    HumanMessage(
+                        content=input_text
+                    )
+                ],
+            },
+            {"recursion_limit": 300},  # Increased recursion limit
+        ):
+            if "__end__" not in s:
+                st.write(s)
+                st.write("---")
+            
+            # Check for timeout
+            if time.time() - start_time > max_execution_time:
+                st.warning("Execution time exceeded. Terminating the process.")
+                break
+    except RetryError as re:
+        st.error(f"Retry error occurred: {re}")
+        logger.error(f"Retry error in super_graph execution: {re}")
+    except ValidationError as ve:
+        st.error(f"Validation error occurred: {ve}")
+        logger.error(f"Validation error in super_graph execution: {ve}")
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+        logger.error(f"Error in super_graph execution: {str(e)}")
+
+if st.button("List Output Files"):
+    files = os.listdir(WORKING_DIRECTORY)
+    if files:
+        st.write("### 📂 Files in working directory:")
+        for file in files:
+            st.write(f"📄 {file}")
+    else:
+        st.write("No files found in the working directory.")
+
+output_files = os.listdir(WORKING_DIRECTORY)
+if output_files:
+    output_file = st.selectbox("Select an output file to download:", output_files)
+
+    if st.button("Download Output Document"):
+        file_path = WORKING_DIRECTORY / output_file
+        if file_path.exists():
+            with file_path.open("rb") as file:
+                st.download_button(
+                    label="📥 Download Output Document",
+                    data=file,
+                    file_name=output_file,
+                )
+        else:
+            st.write("Output document not found.")
+else:
+    st.write("No output files available for download.")
+
+# Cleanup
+if st.button("Clear Working Directory"):
+    for file in WORKING_DIRECTORY.iterdir():
+        if file.is_file():
+            file.unlink()
+    st.success("Working directory cleared.")