Update app.py

app.py

@@ -15,7 +15,84 @@ def main():
     st.header("Info Assistant :" ":books:")
     
     st.markdown("###### Get support of "Info Assistant" , who has in memory a lot of Data Science related articles, if it can't answer based on it's knowledge base, information will be found on the internet:" ":books:")
+
+
+    if "messages" not in st.session_state:
+        st.session_state["messages"] = [
+        {"role": "assistant", "content": "Hi, I'm a chatbot who is  based on respublic of Lithuania law documents. How can I help you?"}
+    ]
+
+
+    search_type = st.selectbox(
+        "Choose search type. Options are [Max marginal relevance search (similarity) , Similarity search (similarity). Default value (similarity)]", 
+        options=["mmr", "similarity"], 
+        index=1  
+    )
+
+    k = st.select_slider(
+        "Select amount of documents to be retrieved. Default value (5): ", 
+        options=list(range(2, 16)), 
+        value=4  
+    )
+    retriever = create_retriever_from_chroma(vectorstore_path="docs/chroma/", search_type=search_type, k=k, chunk_size=350, chunk_overlap=30)
+
+
+
+    # Graph
+    workflow = StateGraph(GraphState)
+
+    # Define the nodes
+    workflow.add_node("ask_question", ask_question)
+    workflow.add_node("retrieve", retrieve)  # retrieve
+    workflow.add_node("grade_documents", grade_documents)  # grade documents
+    workflow.add_node("generate", generate)  # generatae
+    workflow.add_node("web_search", web_search)  # web search
+    workflow.add_node("transform_query", transform_query)
+
+
+    # Build graph
+    workflow.set_entry_point("ask_question")
+    workflow.add_conditional_edges(
+        "ask_question",
+        grade_question_toxicity,
+    
+        {
+        "good": "retrieve",
+        'bad': END,
+        
+        },
+    )
+
+    workflow.add_edge("retrieve", "grade_documents")
+    workflow.add_conditional_edges(
+        "grade_documents",
+        decide_to_generate,
+        {
+            "search": "web_search",
+            "generate": "generate",
+        
+        },
+    )
+    workflow.add_edge("web_search", "generate")
+    workflow.add_conditional_edges(
+        "generate",
+        grade_generation_v_documents_and_question,
+        {
+            "not supported": "generate",
+            "useful": END,
+            "not useful": "transform_query",
+        },
+    )
+
+    workflow.add_edge("transform_query", "retrieve")
+
+    custom_graph = workflow.compile()
+    
+    
+
     
+    if user_question := st.text_input("Ask a question about your documents:"):
+    handle_userinput(user_question,retriever,rag_chain)