Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,178 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+import re
+import json
 import streamlit as st
+from pathlib import Path
+from typing import List, Annotated, Any
+import operator
+import pandas as pd
+from tqdm import tqdm
+from pydantic import BaseModel
+from langchain.embeddings.cohere import CohereEmbeddings
+from langchain_cohere import ChatCohere
+from langchain.document_loaders import DirectoryLoader, TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+import cohere
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage
+from langgraph.graph import StateGraph, START, END, add_messages
+from langgraph.constants import Send
+from langgraph.checkpoint.memory import MemorySaver
+
+load_dotenv()
+os.environ["user_agent"] = "langchain-app/1.0"
+COHERE_API_KEY = os.environ["COHERE_API_KEY"]
+co = cohere.Client(COHERE_API_KEY)
+persist_dir = "./chroma_store"
+
+def prepare_vectorstore():
+    loader = DirectoryLoader("./documents", glob="**/*.txt", loader_cls=TextLoader)
+    documents = loader.load()
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
+    docs = splitter.split_documents(documents)
+    embedding = CohereEmbeddings(
+        model="embed-multilingual-light-v3.0",
+        user_agent="langchain-app/1.0",
+        cohere_api_key=""
+    )
+    vectorstore = Chroma.from_documents(
+        documents=tqdm(docs, desc="Embedding"),
+        embedding=embedding,
+        persist_directory=persist_dir
+    )
+    vectorstore.persist()
+    return vectorstore
+
+if not os.path.exists(persist_dir):
+    prepare_vectorstore()
+
+class State(BaseModel):
+    state: List[str] = []
+    messages: Annotated[list[AnyMessage], add_messages]
+    topic: List[str] = []
+    context: List[str] = []
+    sub_topic_list: List[str] = []
+    sub_topics: Annotated[list[AnyMessage], add_messages]
+    stories: Annotated[list[AnyMessage], add_messages]
+    stories_lst: Annotated[list, operator.add]
+
+class StoryState(BaseModel):
+    retrieved_docs: List[Any] = []
+    stories: Annotated[list[AnyMessage], add_messages]
+    reranked_docs: List[str] = []
+    story_topic: str = ""
+    stories_lst: Annotated[list, operator.add]
+
+def extract_topics(messages):
+    topics = []
+    for message in messages:
+        topics.extend(re.findall(r'- \*\*(.*?)\*\*', message.content))
+    return topics
+
+embedding_llm = CohereEmbeddings(
+    model="embed-multilingual-light-v3.0",
+    user_agent="langchain-app/1.0",
+    cohere_api_key=COHERE_API_KEY
+)
+llm = ChatCohere(
+    api_version="2024-02-15-preview",
+    temperature=0.7,
+    model="command-r-plus-08-2024",
+    cohere_api_key=COHERE_API_KEY
+)
+
+beginner_topic_sys_msg = SystemMessage(content="Suppose you are a middle grader who wants to learn constantly about new topics to get a good score in exams.")
+middle_topic_sys_msg = SystemMessage(content="Suppose you are a college student who wants to learn constantly about new topics to get a good score in exams.")
+advanced_topic_sys_msg = SystemMessage(content="Suppose you are a teacher who wants to learn constantly about new topics to teach your students.")
+
+def retrieve_node(state):
+    topic = state.story_topic
+    query = f"information about {topic}"
+    retriever = Chroma(persist_directory=persist_dir, embedding_function=embedding_llm).as_retriever(search_kwargs={"k": 20})
+    docs = retriever.get_relevant_documents(query)
+    return {"retrieved_docs": docs, "question": query}
+
+def rerank_node(state):
+    topic = state.story_topic
+    query = f"Rerank documents based on how good they explain the topic {topic}"
+    docs = state.retrieved_docs
+    texts = [doc.page_content for doc in docs]
+    rerank_results = co.rerank(query=query, documents=texts, top_n=5, model="rerank-v3.5")
+    top_docs = [texts[result.index] for result in rerank_results.results]
+    return {"reranked_docs": top_docs, "question": query}
+
+def generate_story_node(state):
+    context = "\n\n".join(state.reranked_docs)
+    topic = state.story_topic
+    system_message = """
+    Suppose You're a Amazing story writter and scientific thinker. 
+    You have written hundreds of story books explaining scientific topic in childlike manner that exven amiddle grader could understand.
+    You add a subtle humor to your story to make it more life like.
+    """
+    prompt = f"""
+    Now Use the following context to generate a simple engaging story that explains {topic} in such a way an middle schooler can understand the {topic}.\n
+    Context:\n{context}\n\n
+    Story:
+    """
+    response = llm.invoke([SystemMessage(system_message), HumanMessage(prompt)])
+    return {"stories": response}
+
+def beginner_topic(state: State):
+    prompt = f"What are the beginner-level topics you can learn about {', '.join(state.topic)} in {', '.join(state.context)}?"
+    sub_topics = [llm.invoke([beginner_topic_sys_msg] + [prompt])]
+    return {"message": sub_topics[0], "sub_topics": sub_topics[0]}
+
+def middle_topic(state: State):
+    prompt = f"What are the middle-level topics you can learn about {', '.join(state.topic)} in {', '.join(state.context)}? Don't include the topics below:\n\n{(state.sub_topics)}"
+    sub_topics = [llm.invoke([middle_topic_sys_msg] + [prompt])]
+    return {"message": sub_topics, "sub_topics": sub_topics}
+
+def advanced_topic(state: State):
+    prompt = f"What are the advanced-level topics you can learn about {', '.join(state.topic)} in {', '.join(state.context)}? Don't include the topics below:\n\n{(state.sub_topics)}"
+    sub_topics = [llm.invoke([advanced_topic_sys_msg] + [prompt])]
+    return {"message": sub_topics, "sub_topics": sub_topics}
+
+def topic_extractor(state: State):
+    return {"sub_topic_list": extract_topics(state.sub_topics)}
+
+def dynamic_topic_edges(state: State):
+    return [Send("story_generator", {"story_topic": topic}) for topic in state.sub_topic_list]
+
+story_builder = StateGraph(StoryState)
+story_builder.add_node("Retrieve", retrieve_node)
+story_builder.add_node("Rerank", rerank_node)
+story_builder.add_node("Generate", generate_story_node)
+story_builder.set_entry_point("Retrieve")
+story_builder.add_edge("Retrieve", "Rerank")
+story_builder.add_edge("Rerank", "Generate")
+story_builder.set_finish_point("Generate")
+
+story_graph = story_builder.compile()
+
+main_builder = StateGraph(State)
+main_builder.add_node("beginner_topic", beginner_topic)
+main_builder.add_node("middle_topic", middle_topic)
+main_builder.add_node("advanced_topic", advanced_topic)
+main_builder.add_node("topic_extractor", topic_extractor)
+main_builder.add_node("story_generator", story_graph)
+main_builder.add_edge(START, "beginner_topic")
+main_builder.add_edge("beginner_topic", "middle_topic")
+main_builder.add_edge("middle_topic", "advanced_topic")
+main_builder.add_edge("advanced_topic", "topic_extractor")
+main_builder.add_conditional_edges("topic_extractor", dynamic_topic_edges, ["story_generator"])
+main_builder.add_edge("story_generator", END)
+
+memory = MemorySaver()
+react_graph = main_builder.compile(checkpointer=memory, interrupt_after=["topic_extractor"])
+
+st.title("LangGraph Topic Story Generator")
+topic = st.text_input("Enter a topic", "Human Evolution")
+context = st.text_input("Enter a context", "Science")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if st.button("Generate Stories"):
+    thread = {"configurable": {"thread_id": "1"}}
+    react_graph.invoke({"topic": [topic], "context": [context]}, thread)
+    react_graph.update_state(thread, {"sub_topic_list": ['Early Hominins', 'Fossil Evidence', "Darwin's Theory of Evolution"]})
+    result = react_graph.invoke(None, thread, stream_mode="values")
+    for story in result["stories"]:
+        st.markdown(story.content)