Update the chatbot to an LangGraph agent, handling four tasks including RAG, summary, translate and email. Update prompts. Update Gradio UI.

app/agent.py

@@ -0,0 +1,123 @@
+from langgraph.graph import StateGraph, END
+from langgraph.graph.state import CompiledStateGraph
+from typing import Dict, Any
+from app.rag_chain import get_qa_chain
+from app.models.model import LLM as llm
+from app.summary_chain import get_summary_chain
+from app.translate_chain import get_translate_chain
+from app.rewrite_chain import get_rewrite_chain
+
+# --- OUTPUT COLLECTOR ---
+output_collector = {
+    "answer": None,
+    "summary": None,
+    "translation": None,
+    "formal_email": None
+}
+
+# --- TOOLS ---
+
+def answer_with_rag_tool(input_text: str, input_lang: str, **kwargs) -> str:
+    chain = get_qa_chain(language=input_lang)
+    result = chain.invoke({"input": input_text})
+    output_collector["answer"] = result
+    return result
+
+def summarize_tool(input_text: str, input_lang: str, **kwargs) -> str:
+    llm, prompt = get_summary_chain(input_lang)
+    result = llm.invoke(prompt.format_messages(input=input_text))
+    output_collector["summary"] = result.content
+    return result
+
+def translate_tool(input_text: str, input_lang: str, target_lang: str) -> str:
+    if not target_lang:
+        target_lang = input_lang  # Default to input language if no target language provided
+    llm, prompt = get_translate_chain(input_lang)
+    result = llm.invoke(prompt.format_messages(input=input_text, target_lang=target_lang))
+    output_collector["translation"] = result.content
+    return result
+
+def rewrite_email_tool(input_text: str, input_lang: str, target_lang: str) -> str:
+    llm, prompt = get_rewrite_chain(input_lang)
+    result = llm.invoke(prompt.format_messages(input=input_text, target_lang=target_lang))
+    output_collector["formal_email"] = result.content
+    return result
+
+# --- LANGGRAPH STATE & NODES ---
+
+class AgentState(Dict[str, Any]):
+    user_input: str
+    input_lang: str
+    target_lang: str
+    answer: str = None
+    summary: str = None
+    translation: str = None
+    formal_email: str = None
+
+def node_answer_by_RAG(state: AgentState) -> AgentState:
+    answer = answer_with_rag_tool(state["user_input"], state["input_lang"])
+    state["answer"] = answer
+    return state
+
+def node_summarize(state: AgentState) -> AgentState:
+    summary = summarize_tool(state["answer"], state["input_lang"])
+    state["summary"] = summary
+    return state
+
+def node_translate(state: AgentState) -> AgentState:
+    # Use summary if exists, else answer
+    text_to_translate = state.get("summary") or state["answer"]
+    translation = translate_tool(text_to_translate, state["input_lang"], state["target_lang"])
+    state["translation"] = translation
+    return state
+
+def node_rewrite(state: AgentState) -> AgentState:
+    # Use translation if exists, else summary, else answer
+    text_to_rewrite = state.get("translation") or state["answer"]
+    formal_email = rewrite_email_tool(text_to_rewrite, state["input_lang"], state["target_lang"])
+    state["formal_email"] = formal_email
+    return state
+
+# --- LANGGRAPH GRAPH DEFINITION ---
+
+graph = StateGraph(AgentState)
+graph.add_node("answer_by_RAG", node_answer_by_RAG)
+graph.add_node("summarize", node_summarize)
+graph.add_node("translate", node_translate)
+graph.add_node("rewrite", node_rewrite)
+
+graph.set_entry_point("answer_by_RAG")
+# Remove all conditional edges and decision functions, just allow direct edges for sequential execution if needed
+# (But for the new execution plan, we do not need to set edges between nodes, as we step through them manually)
+# Only keep END edge for completeness
+
+# Add END edges for all possible last nodes
+for node in ["answer_by_RAG", "summarize", "translate", "rewrite"]:
+    graph.add_edge(node, END)
+
+compiled_graph = graph.compile()
+
+# --- MAIN API CALL ---
+
+def run_agent(user_input: str, input_lang: str = "Deutsch", target_lang: str = None, do_summarize: bool = False, do_translate: bool = False, do_email: bool = False):
+    print(user_input, input_lang, target_lang, do_summarize, do_translate, do_email)
+    for k in output_collector:
+        output_collector[k] = None
+
+    state = {
+        "user_input": user_input,
+        "input_lang": input_lang,
+        "target_lang": target_lang
+    }
+    execution_plan = ["answer_by_RAG"]
+    if do_summarize:
+        execution_plan.append("summarize")
+    if do_translate:
+        execution_plan.append("translate")
+    if do_email:
+        execution_plan.append("rewrite")
+
+    for node in execution_plan:
+        state = globals()[f"node_{node}"](state)
+        # Yield the current outputs after each node
+        yield {k: v for k, v in output_collector.items() if v is not None}

app/common_languages.json

@@ -0,0 +1,17 @@
+[
+    "English",
+    "Chinese",
+    "German",
+    "French",
+    "Spanish",
+    "Italian",
+    "Russian",
+    "Japanese",
+    "Korean",
+    "Portuguese",
+    "Arabic",
+    "Hindi",
+    "Dutch",
+    "Turkish",
+    "Polish"
+]

app/config.py

@@ -1 +1,2 @@
-language = "Deutsch" # "English"
+language = "Deutsch" # "English"
+model = "OpenAI" # "AzureOpenAI"

app/gradio_ui.py

@@ -1,6 +1,12 @@
 import gradio as gr
 from app.rag_chain import get_qa_chain
 from app import config
+from app import agent
+import json
+
+# Load common languages for translation dropdown
+with open("app/common_languages.json", "r", encoding="utf-8") as f:
+    common_languages = json.load(f)
 
 def load_example_questions(path="examples/questions.txt"):
     with open(path, "r", encoding="utf-8") as f:
@@ -13,18 +19,6 @@ example_questions = {
     "English": load_example_questions("examples/questions_en.txt"),
 }
 
-def get_chain(language):
-    return get_qa_chain(language=language)
-
-def ask_question(query, language):
-    qa_chain = get_chain(language)
-    result = qa_chain.invoke({"input": query})
-    answer = result["answer"]
-    sources = "\n\nSources / Quellen:\n\n" + "\n".join(
-        [f"{doc.metadata['source']} {doc.metadata.get('article_number', '')}" for doc in result.get("context", [])]
-    )
-    return answer + sources
-
 def from_example_dropdown(selected_example):
     return selected_example
 
@@ -34,80 +28,193 @@ def update_examples(language):
 def get_texts(lang):
     if lang == "Deutsch":
         return {
-            "title": "## LLM-basiert Legal RAG Demo (basierend auf dem Bundesgesetz über die Ausländerinnen und Ausländer und über die Integration)",
+            "title": "# LLM-basiert Legal RAG Agent Demo (basierend auf dem Bundesgesetz über die Ausländerinnen und Ausländer und über die Integration)",
             "input_text": "Fragen Sie etwas zum Ausländer- und Integrationsgesetz ...",
             "submit_button": "Absenden",
-            "output_text": "Antwort mit Quellen",
+            "output_text": "Antwort",
             "lang_label": "🌐 Sprache",
             "example_label": "Beispiel-Frage auswählen",
+            "agent_tab_label": "Agent Spielplatz",
+            "agent_input_label": "Frage an den Agenten...",
+            "agent_submit_button": "Agent fragen",
+            "agent_output_label": "Antwort des Agenten",
+            "agent_answer_label": "Agent Antwort",
+            "agent_summary_label": "Agent Zusammenfassung",
+            "agent_translation_label": "Agent Übersetzung",
+            "agent_email_label": "Agent formelle E-Mail",
+            "sources_label": "Quellen",
+            "summarize_btn": "Antwort zusammenfassen",
+            "translate_btn": "Antwort übersetzen in [NEUE SPRACHE]",
+            "email_btn": "E-Mail an Mandanten schreiben"
         }
     else:
         return {
-            "title": "## LLM-Powered Legal RAG Demo (based on Federal Act on Foreign Nationals and Integration)",
+            "title": "# LLM-Powered Legal RAG Agent Demo (based on Federal Act on Foreign Nationals and Integration)",
             "input_text": "Ask something about the Federal Act on Foreign Nationals and Integration ...",
             "submit_button": "Submit",
-            "output_text": "Answer with sources",
+            "output_text": "Answer",
             "lang_label": "🌐 Language",
             "example_label": "Choose an example question",
+            "agent_tab_label": "Agent Playground",
+            "agent_input_label": "Ask the agent something...",
+            "agent_submit_button": "Ask Agent",
+            "agent_output_label": "Agent's Response",
+            "agent_answer_label": "Agent Answer",
+            "agent_summary_label": "Agent Summary",
+            "agent_translation_label": "Agent Translation",
+            "agent_email_label": "Agent Formal Email",
+            "sources_label": "Sources",
+            "summarize_btn": "Summarize the answer",
+            "translate_btn": "Translate the answer to [NEW LANGUAGE]",
+            "email_btn": "Write an email to the client"
         }
 
 with gr.Blocks() as iface:
+    
     state = gr.State(config.language)
     txt = get_texts(config.language)
-    title = gr.Markdown(txt["title"])
+    title = gr.Markdown(txt["title"], elem_id="title_elem")
+
+    # Language Dropdown
+    with gr.Row():
+        lang_dropdown = gr.Dropdown(
+            label=txt["lang_label"],
+            choices=["Deutsch", "English"],
+            value=config.language,
+            elem_id="lang_dropdown_elem"
+        )
 
+    # Agent Playground Content (no Tab)
+    agent_title_md = gr.Markdown("## " + txt["agent_tab_label"])
     with gr.Row():
         with gr.Column(scale=1):
-            lang_dropdown = gr.Dropdown(
-                label=txt["lang_label"],
-                choices=["Deutsch", "English"],
-                value=config.language
+            agent_input_box = gr.Textbox(
+                label=txt["agent_input_label"],
+                lines=5,
+                elem_id="agent_input_box_elem"
             )
-            input_box = gr.Textbox(
-                label=txt["input_text"],
-                lines=2
-            )
-            example_dropdown = gr.Dropdown(
+            agent_example_dropdown = gr.Dropdown(
                 label=txt["example_label"],
                 choices=[''] + example_questions[config.language],
-                value=''
+                value='',
+                elem_id="agent_example_dropdown_elem"
             )
-            example_dropdown.change(fn=from_example_dropdown, inputs=example_dropdown, outputs=input_box)
-            lang_dropdown.change(fn=update_examples, inputs=lang_dropdown, outputs=example_dropdown)
-            submit_btn = gr.Button(txt["submit_button"])
-
+            agent_example_dropdown.change(fn=from_example_dropdown, inputs=agent_example_dropdown, outputs=agent_input_box)
+            # Function checkboxes
+            summarize_chk = gr.Checkbox(label=txt["summarize_btn"], value=False, elem_id="summarize_chk")
+            translate_chk = gr.Checkbox(label=txt["translate_btn"], value=False, elem_id="translate_chk")
+            email_chk = gr.Checkbox(label=txt["email_btn"], value=False, elem_id="email_chk")
+            # Language dropdown for translation (hidden unless translate_chk is True)
+            translate_lang_dropdown = gr.Dropdown(
+                label="Target Language", choices=common_languages, value="English", visible=False, elem_id="translate_lang_dropdown"
+            )
+            def show_translate_dropdown(checked):
+                return gr.update(visible=checked)
+            translate_chk.change(fn=show_translate_dropdown, inputs=translate_chk, outputs=translate_lang_dropdown)
+            agent_submit_btn = gr.Button(txt["agent_submit_button"], elem_id="agent_submit_btn_elem", variant="primary", elem_classes=["agent-submit-btn"])
+            # Add custom CSS for button colors and equal width
+            gr.HTML("""
+            <style>
+            .agent-submit-btn {
+                background-color: #bbdefb !important; /* light blue */
+                color: #333 !important;
+                border: 1px solid #64b5f6 !important;
+                font-weight: 600;
+                font-size: 1.05em;
+            }
+            .agent-func-btn {
+                background-color: #fff9c4 !important; /* light yellow */
+                color: #333 !important;
+                border: 1px solid #f7e48b !important;
+                min-width: 180px;
+                width: 100%;
+                box-sizing: border-box;
+                font-weight: 600;
+                font-size: 1.05em;
+                text-align: center !important;
+                white-space: normal !important;
+            }
+            </style>
+            """)
+            
+            agent_sources_box = gr.Textbox(label=txt["sources_label"], lines=4, elem_id="agent_sources_box_elem")
         with gr.Column(scale=1):
-            output_box = gr.Textbox(label=txt["output_text"], lines=12)
+            agent_answer_box = gr.Textbox(label=txt["agent_answer_label"], lines=6, elem_id="agent_answer_box_elem")
+            agent_summary_box = gr.Textbox(label=txt["agent_summary_label"], lines=3, elem_id="agent_summary_box_elem")
+            agent_translation_box = gr.Textbox(label=txt["agent_translation_label"], lines=3, elem_id="agent_translation_box_elem")
+            agent_email_box = gr.Textbox(label=txt["agent_email_label"], lines=3, elem_id="agent_email_box_elem")
 
-        
+    # Main submit: answer, summary, translation, email, sources
+    def ask_agent_with_sources_stream(query, language, do_summarize, do_translate, do_email, translate_target_lang):
+        stream = agent.run_agent(
+            query,
+            language,
+            target_lang=translate_target_lang if do_translate else None,
+            do_summarize=do_summarize,
+            do_translate=do_translate,
+            do_email=do_email
+        )
+        # Initial empty values
+        answer = summary = translation = formal_email = sources = ""
+        for result in stream:
+            print(result)
+            answer_dict = result.get("answer", dict())
+            answer = answer_dict.get("answer", "")
+            context = answer_dict.get("context", [])
+            sources = "\n".join([
+                f"{doc.metadata['source']} {doc.metadata.get('article_number', '')}" for doc in context
+            ]) if context else ""
 
-    submit_btn.click(
-        fn=ask_question,
-        inputs=[input_box, lang_dropdown],
-        outputs=output_box
+            summary = result.get("summary", summary)
+            translation = result.get("translation", translation)
+            formal_email = result.get("formal_email", formal_email)
+            # (sources extraction logic as before)
+            # yield the current state of all outputs
+            yield answer or "", summary or "", translation or "", formal_email or "", sources or ""
+
+    agent_submit_btn.click(
+        fn=ask_agent_with_sources_stream,
+        inputs=[agent_input_box, lang_dropdown, summarize_chk, translate_chk, email_chk, translate_lang_dropdown],
+        outputs=[agent_answer_box, agent_summary_box, agent_translation_box, agent_email_box, agent_sources_box]
     )
 
-    def update_labels(lang):
+    def update_all_labels_and_ui(lang):
         txt = get_texts(lang)
-        return (
+        return [
             gr.update(value=txt["title"]),
-            gr.update(label=txt["input_text"]),
-            gr.update(label=txt["output_text"]),
-            gr.update(value=txt["submit_button"]),
-            gr.update(label=txt["lang_label"]),
-            gr.update(label=txt["example_label"], choices=[''] + example_questions[lang], value='')
-        )
+            gr.update(label=txt["lang_label"], value=lang),
+            gr.update(label=txt["example_label"], choices=[''] + example_questions[lang], value=''),
+            gr.update(value="## " + txt["agent_tab_label"]),
+            gr.update(label=txt["agent_input_label"]),
+            gr.update(value=txt["agent_submit_button"]),
+            gr.update(label=txt["sources_label"]),
+            gr.update(label=txt["agent_answer_label"]),
+            gr.update(label=txt["agent_summary_label"]),
+            gr.update(label=txt["agent_translation_label"]),
+            gr.update(label=txt["agent_email_label"]),
+            gr.update(label=txt["summarize_btn"]),  # summarize_chk
+            gr.update(label=txt["translate_btn"]),  # translate_chk
+            gr.update(label=txt["email_btn"]),      # email_chk
+        ]
 
     lang_dropdown.change(
-        fn=update_labels,
+        fn=update_all_labels_and_ui,
         inputs=lang_dropdown,
         outputs=[
             title,
-            input_box,
-            output_box,
-            submit_btn,
             lang_dropdown,
-            example_dropdown,
+            agent_example_dropdown,
+            agent_title_md,
+            agent_input_box,
+            agent_submit_btn,
+            agent_sources_box,
+            agent_answer_box,
+            agent_summary_box,
+            agent_translation_box,
+            agent_email_box,
+            summarize_chk,
+            translate_chk,
+            email_chk
         ]
     )
 

app/models/model.py

@@ -1,11 +1,12 @@
 from langchain_huggingface import HuggingFaceEmbeddings
 from app.models.AzureOpenAI import azure_llm
-
+from app.models.OpenAI import llm as openai_llm
+from app.config import model
 
 # model = TinyLlamaModel()
 # LLM = 
 
-LLM = azure_llm
+LLM = openai_llm if model == "OpenAI" else azure_llm
 
 Embedding_model_en = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # only English
 Embedding_model_de = HuggingFaceEmbeddings(model_name="danielheinz/e5-base-sts-en-de")

app/rewrite_chain.py

@@ -0,0 +1,12 @@
+from langchain_core.prompts import ChatPromptTemplate
+from app.models.model import LLM
+from app.utils.prompts import rewrite_prompts
+
+def get_rewrite_chain(language: str = "English"):
+    system_prompt = rewrite_prompts[language]
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        ("human", "{input}"),
+    ])
+    
+    return LLM, prompt

app/summary_chain.py

@@ -0,0 +1,14 @@
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.prompts import ChatPromptTemplate
+from app.models.model import LLM
+from app.managers import vector_manager as vm
+from app.utils.prompts import summary_prompts
+
+def get_summary_chain(language: str = "Deutsch"):
+    system_prompt = summary_prompts[language]
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        ("human", "{input}"),
+    ])
+    
+    return LLM, prompt

app/translate_chain.py

@@ -0,0 +1,12 @@
+from langchain_core.prompts import ChatPromptTemplate
+from app.models.model import LLM
+from app.utils.prompts import translate_prompts
+
+def get_translate_chain(language: str = "English"):
+    system_prompt = translate_prompts[language]
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        ("human", "{input}"),
+    ])
+    
+    return LLM, prompt

app/utils/prompts.py

@@ -1,4 +1,3 @@
-
 prompts = {
     'English':
     """
@@ -39,4 +38,43 @@ QUESTION:
     """
 
 
+}
+
+summary_prompts = {
+    'English':
+    """
+You are a legal assistant. Summarize the following answer as concisely as possible, keeping all key legal points and references. Use clear, formal language.
+
+""",
+    'Deutsch':
+    """
+Sie sind ein juristischer Assistent. Fassen Sie die folgende Antwort so knapp wie möglich zusammen und behalten Sie alle wichtigen rechtlichen Punkte und Verweise bei. Verwenden Sie eine klare, formelle Sprache.
+
+"""
+}
+
+translate_prompts = {
+    'English':
+    """
+You are a legal translator. Translate the following answer into {target_lang}, preserving all legal terminology and references. Be accurate and formal.
+
+""",
+    'Deutsch':
+    """
+Sie sind ein juristischer Übersetzer. Übersetzen Sie die folgende Antwort ins {target_lang} und erhalten Sie dabei alle juristischen Begriffe und Verweise. Seien Sie genau und formell.
+
+"""
+}
+
+rewrite_prompts = {
+    'English':
+    """
+You are a legal assistant. Rewrite the following answer as a formal email to a client, using a professional and courteous tone. Begin with a greeting, present the answer, and end with a closing. Keep the language to {target_lang}. 
+
+""",
+    'Deutsch':
+    """
+Sie sind ein juristischer Assistent. Formulieren Sie die folgende Antwort als formelle E-Mail an einen Mandanten um. Verwenden Sie einen professionellen und höflichen Ton. Beginnen Sie mit einer Anrede, präsentieren Sie die Antwort und schließen Sie mit einem Gruß ab. Behalten Sie die Originalsprache als {target_lang}.
+
+"""
 }

requirements.txt

@@ -7,6 +7,7 @@ gradio
 langchain-community
 langchain_huggingface
 langchain_openai
+langgraph
 pydantic
 ipykernel
 pypdf