Create interim_v1.py

interim_v1.py

@@ -0,0 +1,164 @@
+import os
+import chromadb
+import streamlit as st
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langchain.agents import AgentExecutor, create_openai_tools_agent
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_experimental.tools import PythonREPLTool
+from langchain_community.document_loaders import DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceBgeEmbeddings
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langgraph.graph import StateGraph, END
+from langchain_core.documents import Document
+from typing import Annotated, Sequence, TypedDict
+import functools
+import operator
+from langchain_core.tools import tool
+
+
+# Clear ChromaDB cache to fix tenant issue
+chromadb.api.client.SharedSystemClient.clear_system_cache()
+
+# Load environment variables
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+
+if not OPENAI_API_KEY or not TAVILY_API_KEY:
+    st.error("Please set OPENAI_API_KEY and TAVILY_API_KEY in your environment variables.")
+    st.stop()
+
+# Initialize API keys and LLM
+llm = ChatOpenAI(model="gpt-4-1106-preview", openai_api_key=OPENAI_API_KEY)
+
+# Utility Functions
+def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        MessagesPlaceholder(variable_name="messages"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ])
+    agent = create_openai_tools_agent(llm, tools, prompt)
+    return AgentExecutor(agent=agent, tools=tools)
+
+def agent_node(state, agent, name):
+    result = agent.invoke(state)
+    return {"messages": [HumanMessage(content=result["output"], name=name)]}
+
+@tool
+def RAG(state):
+    """Use this tool to execute RAG. If the question is related to Japan or Sports, this tool retrieves the results."""
+    st.session_state.outputs.append('-> Calling RAG ->')
+    question = state
+    template = """Answer the question based only on the following context:\n{context}\nQuestion: {question}"""
+    prompt = ChatPromptTemplate.from_template(template)
+    retrieval_chain = (
+        {"context": retriever, "question": RunnablePassthrough()} |
+        prompt |
+        llm |
+        StrOutputParser()
+    )
+    result = retrieval_chain.invoke(question)
+    return result
+
+# Load Tools and Retriever
+tavily_tool = TavilySearchResults(max_results=5, tavily_api_key=TAVILY_API_KEY)
+python_repl_tool = PythonREPLTool()
+
+# File Upload Section
+st.title("Multi-Agent Workflow Demonstration")
+uploaded_files = st.file_uploader("Upload your source files (TXT)", accept_multiple_files=True, type=['txt'])
+
+if uploaded_files:
+    docs = []
+    for uploaded_file in uploaded_files:
+        content = uploaded_file.read().decode("utf-8")
+        docs.append(Document(page_content=content, metadata={"name": uploaded_file.name}))
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10, length_function=len)
+    new_docs = text_splitter.split_documents(documents=docs)
+    embeddings = HuggingFaceBgeEmbeddings(model_name="BAAI/bge-base-en-v1.5", model_kwargs={'device': 'cpu'}, encode_kwargs={'normalize_embeddings': True})
+    db = Chroma.from_documents(new_docs, embeddings)
+    retriever = db.as_retriever(search_kwargs={"k": 4})
+else:
+    retriever = None
+    st.warning("Please upload at least one text file to proceed.")
+    st.stop()
+
+# Create Agents
+research_agent = create_agent(llm, [tavily_tool], "You are a web researcher.")
+code_agent = create_agent(llm, [python_repl_tool], "You may generate safe python code to analyze data and generate charts using matplotlib.")
+RAG_agent = create_agent(llm, [RAG], "Use this tool when questions are related to Japan or Sports category.")
+
+research_node = functools.partial(agent_node, agent=research_agent, name="Researcher")
+code_node = functools.partial(agent_node, agent=code_agent, name="Coder")
+rag_node = functools.partial(agent_node, agent=RAG_agent, name="RAG")
+
+members = ["RAG", "Researcher", "Coder"]
+system_prompt = (
+    "You are a supervisor managing these workers: {members}. Respond with the next worker or FINISH. "
+    "Use RAG tool for Japan or Sports questions."
+)
+options = ["FINISH"] + members
+function_def = {
+    "name": "route", "description": "Select the next role.",
+    "parameters": {
+        "title": "routeSchema", "type": "object",
+        "properties": {"next": {"anyOf": [{"enum": options}]}}, "required": ["next"]
+    }
+}
+prompt = ChatPromptTemplate.from_messages([
+    ("system", system_prompt),
+    MessagesPlaceholder(variable_name="messages"),
+    ("system", "Given the conversation above, who should act next? Select one of: {options}"),
+]).partial(options=str(options), members=", ".join(members))
+
+supervisor_chain = (prompt | llm.bind_functions(functions=[function_def], function_call="route") | JsonOutputFunctionsParser())
+
+# Build Workflow
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    next: str
+
+workflow = StateGraph(AgentState)
+workflow.add_node("Researcher", research_node)
+workflow.add_node("Coder", code_node)
+workflow.add_node("RAG", rag_node)
+workflow.add_node("supervisor", supervisor_chain)
+
+for member in members:
+    workflow.add_edge(member, "supervisor")
+conditional_map = {k: k for k in members}
+conditional_map["FINISH"] = END
+workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)
+workflow.set_entry_point("supervisor")
+graph = workflow.compile()
+
+# Streamlit UI
+if 'outputs' not in st.session_state:
+    st.session_state.outputs = []
+
+user_input = st.text_area("Enter your task or question:")
+
+def run_workflow(task):
+    st.session_state.outputs.clear()
+    st.session_state.outputs.append(f"User Input: {task}")
+    for state in graph.stream({"messages": [HumanMessage(content=task)]}):
+        if "__end__" not in state:
+            st.session_state.outputs.append(str(state))
+            st.session_state.outputs.append("----")
+
+if st.button("Run Workflow"):
+    if user_input:
+        run_workflow(user_input)
+    else:
+        st.warning("Please enter a task or question.")
+
+st.subheader("Workflow Output:")
+for output in st.session_state.outputs:
+    st.text(output)