added gradio

agent/graph.py

@@ -6,31 +6,46 @@ from functools import partial
 from .state import AgentState
 from .nodes import *
 
-def create_agent_graph(llm, tools) -> StateGraph:
+
+def create_agent_graph(llm, vectordb, search_tools) -> StateGraph:
     """Create the RAG agent graph."""
-    llm_with_tools = llm.with_tools(tools)
 
     graph = StateGraph(AgentState)
     checkpointer = MemorySaver()
 
-    tool_node = ToolNode(tools)
+    web_search_tool_node = ToolNode(search_tools) 
+
+    # --- Nodes ---
+    graph.add_node("router_node", partial(router_node, llm=llm)) 
+    graph.add_node("vectordb_node", partial(vectordb_node, vectorstore=vectordb))
+    graph.add_node("web_search_node", web_search_tool_node)
+    graph.add_node("generate_node", partial(generate_node, llm=llm))
+
+    # --- Edges ---
+    graph.add_edge(START, "router_node")
+
+    graph.add_conditional_edges(
+        "router_node",
+        routing_logic,
+        {
+            # Output from routing_logic -> Target Node Name
+            "vectordb_node": "vectordb_node",
+            "web_search_node": "web_search_node",
+            "generate_node": "generate_node", 
+            # If your logic has an 'else' that returns END, you don't list it here.
+        }
+    )
 
-    # Nodes
-    graph.add_node("summarize", partial(analyze_chat_and_summarize, llm=llm)) # summarize last 6 messages
-    graph.add_node("analyze_rewrite", partial(analyze_and_rewrite_query, llm=llm)) # analyze and rewrite query
-    graph.add_node("agent", partial(agent_node, llm_with_tools=llm_with_tools)) # generate answer based on retrieved info
-    graph.add_node("tools", tool_node)
+    graph.add_edge("vectordb_node", "generate_node")
+    graph.add_edge("web_search_node", "generate_node")
 
-    graph.add_edge(START, "summarize")
-    graph.add_edge("summarize", "analyze_rewrite")
-    graph.add_conditional_edges("analyze_rewrite", route_after_rewrite)
-    graph.add_edge("human_input", "analyze_rewrite")
-    graph.add_conditional_edges("agent", tools_condition)
-    graph.add_edge("tools", "agent")
+    graph.add_edge("generate_node", END)
 
     agent_graph = graph.compile(
         checkpointer=checkpointer,
-        interrupt_before=["human_input"]
     )
 
-    return agent_graph
+    return agent_graph
+
+if __name__ == "__main__":
+    pass

agent/more_nodes.py

@@ -0,0 +1,97 @@
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, RemoveMessage
+from typing import Literal
+
+from .state import AgentState, QueryAnalysis
+from .prompts import *
+
+def analyze_chat_and_summarize(state: AgentState, llm):
+    """
+    Analyzes chat history and summarizes key points for context.
+    """
+    if len(state["messages"]) < 4:  # Need some history to summarize
+        return {"conversation_summary": ""}
+
+    # Extract relevant messages (excluding current query and system messages)
+    relevant_msgs = [
+        msg for msg in state["messages"][:-1]  # Exclude current query
+        if isinstance(msg, (HumanMessage, AIMessage))
+        and not getattr(msg, "tool_calls", None)
+    ]
+
+    if not relevant_msgs:
+        return {"conversation_summary": ""}
+
+    summary_prompt = """**Summarize the key topics and context from this conversation concisely (1-2 sentences max).**
+    Discard irrelevant information, such as misunderstandings or off-topic queries/responses.
+    If there are no key topics, return an empty string.
+
+    """
+    for msg in relevant_msgs[-6:]:  # Last 6 messages for context
+        role = "User" if isinstance(msg, HumanMessage) else "Assistant"
+        summary_prompt += f"{role}: {msg.content}\n"
+
+    summary_prompt += "\nBrief Summary:"
+    summary_response = llm.with_config(temperature=0.3).invoke([SystemMessage(content=summary_prompt)])
+    return {"conversation_summary": summary_response.content}
+
+def analyze_and_rewrite_query(state: AgentState, llm):
+    """
+    Analyzes user query and rewrites it for clarity, optionally using conversation context.
+    """
+    last_message = state["messages"][-1]
+    conversation_summary = state.get("conversation_summary", "")
+
+    context_section = (
+        f"**Conversation Context:**\n{conversation_summary}"
+        if conversation_summary.strip()
+        else "**Conversation Context:**\n[First query in conversation]"
+    )
+
+    # Create analysis prompt
+    query_analysis_prompt = get_query_analysis_prompt(last_message.content, conversation_summary)
+
+    llm_with_structure = llm.with_config(temperature=0.3).with_structured_output(QueryAnalysis)
+    response = llm_with_structure.invoke([SystemMessage(content=query_analysis_prompt)])
+
+    if response.is_clear:
+        # Remove all non-system messages
+        delete_all = [
+            RemoveMessage(id=m.id)
+            for m in state["messages"]
+            if not isinstance(m, SystemMessage)
+        ]
+
+        # Format rewritten query
+        rewritten = (
+            "\n".join([f"{i+1}. {q}" for i, q in enumerate(response.questions)])
+            if len(response.questions) > 1
+            else response.questions[0]
+        )
+        return {
+            "questionIsClear": True,
+            "messages": delete_all + [HumanMessage(content=rewritten)]
+        }
+    else:
+        clarification = response.clarification_needed or "I need more information to understand your question."
+        return {
+            "questionIsClear": False,
+            "messages": [AIMessage(content=clarification)]
+        }
+
+def human_input_node(state: AgentState):
+    """Placeholder node for human-in-the-loop interruption"""
+    return {}
+
+def route_after_rewrite(state: AgentState) -> Literal["agent", "human_input"]:
+    """Route to agent if question is clear, otherwise wait for human input"""
+    return "agent" if state.get("questionIsClear", False) else "human_input"
+
+def agent_node(state: AgentState, llm_with_tools):
+    """Main agent node that processes queries using tools"""
+    system_prompt = get_system_prompt()
+    messages = [system_prompt] + state["messages"]
+    response = llm_with_tools.invoke(messages)
+    return {"messages": [response]}
+
+if __name__ == "__main__":
+    pass

agent/nodes.py

@@ -1,97 +1,61 @@
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, RemoveMessage
 from typing import Literal
+from langgraph.graph import START, END
 
 from .state import AgentState, QueryAnalysis
 from .prompts import *
+from .tools import intialize_chroma_vectorstore
 
-def analyze_chat_and_summarize(state: AgentState, llm):
-    """
-    Analyzes chat history and summarizes key points for context.
-    """
-    if len(state["messages"]) < 4:  # Need some history to summarize
-        return {"conversation_summary": ""}
-
-    # Extract relevant messages (excluding current query and system messages)
-    relevant_msgs = [
-        msg for msg in state["messages"][:-1]  # Exclude current query
-        if isinstance(msg, (HumanMessage, AIMessage))
-        and not getattr(msg, "tool_calls", None)
-    ]
-
-    if not relevant_msgs:
-        return {"conversation_summary": ""}
-
-    summary_prompt = """**Summarize the key topics and context from this conversation concisely (1-2 sentences max).**
-    Discard irrelevant information, such as misunderstandings or off-topic queries/responses.
-    If there are no key topics, return an empty string.
 
+def router_node(state: AgentState, llm):
     """
-    for msg in relevant_msgs[-6:]:  # Last 6 messages for context
-        role = "User" if isinstance(msg, HumanMessage) else "Assistant"
-        summary_prompt += f"{role}: {msg.content}\n"
-
-    summary_prompt += "\nBrief Summary:"
-    summary_response = llm.with_config(temperature=0.3).invoke([SystemMessage(content=summary_prompt)])
-    return {"conversation_summary": summary_response.content}
+    Takes the query (and history). Decides the next step: vectordb, tools, or refuse.
+    """
+    query = state["messages"][-1].content
+    rag_method_prompt = determine_rag_method_prompt()
+    rag_method_result = llm.invoke([rag_method_prompt, HumanMessage(content=query)])
+    rag_method = rag_method_result.content.strip().upper()
+    state["rag_method"] = rag_method 
+    return state
+
+def routing_logic(self, state: AgentState) -> str:
+    rag_method = state["rag_method"]
+    if rag_method == "RAG":
+        return "vectordb_node"
+    elif rag_method == "WEBSEARCH":
+        return "web_search_node"
+    elif rag_method == "GENERAL":
+        return "generate_node" # fallback to generate_node if the question do not requires RAG or websearch
+    else:
+        # If the LLM violates the prompt and outputs an unknown word,
+        print(f"ERROR: Router returned unclassified intent: {rag_method}. Terminating flow.")
+        return END
 
-def analyze_and_rewrite_query(state: AgentState, llm):
+def vectordb_node(state: AgentState, llm, vectorstore):
     """
-    Analyzes user query and rewrites it for clarity, optionally using conversation context.
+    Use vectordb to answer the query.
     """
-    last_message = state["messages"][-1]
-    conversation_summary = state.get("conversation_summary", "")
-
-    context_section = (
-        f"**Conversation Context:**\n{conversation_summary}"
-        if conversation_summary.strip()
-        else "**Conversation Context:**\n[First query in conversation]"
+    context_docs = vectorstore.similarity_search(
+        query=state["messages"][-1].content,
+        k=5
     )
-
-    # Create analysis prompt
-    query_analysis_prompt = get_query_analysis_prompt(last_message.content, conversation_summary)
-
-    llm_with_structure = llm.with_config(temperature=0.3).with_structured_output(QueryAnalysis)
-    response = llm_with_structure.invoke([SystemMessage(content=query_analysis_prompt)])
-
-    if response.is_clear:
-        # Remove all non-system messages
-        delete_all = [
-            RemoveMessage(id=m.id)
-            for m in state["messages"]
-            if not isinstance(m, SystemMessage)
-        ]
-
-        # Format rewritten query
-        rewritten = (
-            "\n".join([f"{i+1}. {q}" for i, q in enumerate(response.questions)])
-            if len(response.questions) > 1
-            else response.questions[0]
-        )
-        return {
-            "questionIsClear": True,
-            "messages": delete_all + [HumanMessage(content=rewritten)]
-        }
-    else:
-        clarification = response.clarification_needed or "I need more information to understand your question."
-        return {
-            "questionIsClear": False,
-            "messages": [AIMessage(content=clarification)]
-        }
-
-def human_input_node(state: AgentState):
-    """Placeholder node for human-in-the-loop interruption"""
-    return {}
-
-def route_after_rewrite(state: AgentState) -> Literal["agent", "human_input"]:
-    """Route to agent if question is clear, otherwise wait for human input"""
-    return "agent" if state.get("questionIsClear", False) else "human_input"
-
-def agent_node(state: AgentState, llm_with_tools):
-    """Main agent node that processes queries using tools"""
-    system_prompt = get_system_prompt()
-    messages = [system_prompt] + state["messages"]
-    response = llm_with_tools.invoke(messages)
-    return {"messages": [response]}
-
+    context = "\n\n".join([doc.page_content for doc in context_docs])
+    state["context"] = context
+    return state
+
+def generate_node(state: AgentState, llm):
+    messages = state["messages"][-10:]  # Limit to last 10 messages to handle token limit
+    context = state.get("context", [])
+    
+    system_content = get_system_prompt()
+
+    if context:
+        system_content += f"\n\nRelevant Context:\n{context}"
+    
+    messages_with_system = [SystemMessage(content=system_content)] + messages
+    response = llm.invoke(messages_with_system)
+
+    return {'messages': [response]}
+    
 if __name__ == "__main__":
     pass

agent/prompts.py

@@ -16,6 +16,19 @@ You are an intelligent assistant that MUST use the available tools to answer que
 7. **Return the final answer** derived from the most relevant results.
 """)
 
+def determine_rag_method_prompt() -> str:
+    return SystemMessage(content="""
+You are an rag method classification model. Given the user's query, you must classify the method to use
+as one and only one of the following options:
+
+1. **RAG**: The query likely relates to the internal, domain-specific documents you have access to.
+2. **WEBSEARCH**: The query requires real-time facts, general knowledge, or external information not in your documents.
+3. **GENERAL**: The query can be answered based on your existing knowledge without external resources.
+
+Respond STRICTLY with only one of these words: RAG, WEBSEARCH, or GENERAL. Do not include any punctuation, explanation, or extra text.
+"""
+    )
+
 def get_conversation_summary_prompt(messages):
     """Generate a prompt for conversation summarization."""
     summary_prompt = """**Summarize the key topics and context from this conversation concisely (1-2 sentences max).**

agent/state.py

@@ -2,13 +2,23 @@ from typing import TypedDict, Annotated, Sequence, Optional, List
 from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage, ToolMessage
 from langgraph.graph.message import add_messages
 from pydantic import BaseModel, Field
+from enum import Enum
 
+class RAG_method(str, Enum):
+    RAG = "RAG"
+    WEBSEARCH = "WEBSEARCH"
+    GENERAL = "GENERAL"
 
 class AgentState(TypedDict):
     messages: Annotated[Sequence[AnyMessage], add_messages]
+    rag_method: RAG_method
+    context: Optional[str]
+
     questionIsClear: bool
     conversation_summary: str = ""
-    
+
+
+
 class QueryAnalysis(BaseModel):
     """Structured output for query analysis"""
     is_clear: bool = Field(description="Indicates if the user's question is clear and answerable")

agent/tools.py

@@ -23,30 +23,28 @@ def intialize_chroma_vectorstore():
     )
     return vectorstore
 
-
 @tool
-def search_chroma(vectorstore: Chroma, query: str, k: int = 5) -> List[dict]:
-    """Search for the top K most relevant chunks from Chroma vector store.
-
+def web_search_tavily(query: str) -> dict:
+    """Search Tavily for a query and return up to 3 results.
     Args:
-        query: Search query string
-        k: Number of results to return
+        query: The search query.
+    Returns:
+        dict with key 'web_results', containing a list of search results with
+        'source', 'page', and 'content'.
     """
     try:
-        results = vectorstore.similarity_search(query, k=k, score_threshold=0.7)
-        
-        return [
+        search_docs = TavilySearchResults(max_results=3).invoke(input=query)
+        results = [
             {
-                "content": doc.page_content,
-                "parent_id": doc.metadata.get("parent_id", ""),
-                "source": doc.metadata.get("source", "")
+                "title": doc.get("title", ""),
+                "url": doc.get("url", ""),
+                "content": doc.get("content", ""),
             }
-            for doc in results
+            for doc in search_docs
         ]
-    
+        return {"web_results": results}
     except Exception as e:
-        print(f"Error searching chunks: {e}")
-        return []
+        return {"web_results": f"Error retrieving results: {str(e)}"}
     
 @tool 
 def wikipedia_search(query: str) -> dict:

config.py

@@ -2,7 +2,11 @@ import os
 
 configs = {
     "DATA_PATH": "./docs/markdowns",
-    "PERSIST_PATH": "./knowledge_base/chroma_data",
+    "PERSIST_PATH": "./chroma_data",
     "EMBEDDING_MODEL_NAME": "sentence-transformers/all-mpnet-base-v2",
     "COLLECTION_NAME": "langchain_mpnet_collection"
 }
+
+if __name__ == "__main__":
+    for key, value in configs.items():
+        os.environ[key] = value

core/chat_interface.py

@@ -0,0 +1,196 @@
+import gradio as gr
+from core.rag_agent import RAGAgent
+from core.document_manager import DocumentManager
+import os
+
+# Initialize components
+doc_manager = DocumentManager()
+rag_agent = None
+
+def initialize_agent():
+    """Initialize RAG agent lazily"""
+    global rag_agent
+    if rag_agent is None:
+        rag_agent = RAGAgent()
+    return rag_agent
+
+def upload_files(files):
+    """Handle file uploads"""
+    if not files:
+        return "No files selected", get_file_list()
+    
+    results = []
+    for file in files:
+        try:
+            result = doc_manager.add_document(file.name)
+            results.append(result)
+        except Exception as e:
+            results.append(f"Error processing {os.path.basename(file.name)}: {str(e)}")
+    
+    return "\n".join(results), get_file_list()
+
+def get_file_list():
+    """Get list of documents in the knowledge base"""
+    try:
+        files = doc_manager.list_documents()
+        if not files:
+            return "No documents in knowledge base"
+        return "\n".join([f"• {f}" for f in files])
+    except Exception as e:
+        return f"Error listing files: {str(e)}"
+
+def clear_database():
+    """Clear all documents from the knowledge base"""
+    try:
+        result = doc_manager.clear_all()
+        return result, get_file_list()
+    except Exception as e:
+        return f"Error clearing database: {str(e)}", get_file_list()
+
+def chat_with_agent(message, history):
+    """Handle chat interactions with the RAG agent"""
+    if not message.strip():
+        return history
+    
+    try:
+        agent = initialize_agent()
+        
+        # Stream the agent's response
+        response_text = ""
+        for event in agent.agent_graph.stream(
+            {"messages": [("user", message)]},
+            agent.get_config(),
+            stream_mode="values"
+        ):
+            if "messages" in event and len(event["messages"]) > 0:
+                last_message = event["messages"][-1]
+                if hasattr(last_message, "content"):
+                    response_text = last_message.content
+        
+        if not response_text:
+            response_text = "I apologize, but I couldn't generate a response. Please try again."
+        
+        return response_text
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+def reset_conversation():
+    """Reset the conversation thread"""
+    global rag_agent
+    if rag_agent:
+        rag_agent.reset_thread()
+    return None  # Clear chat history
+
+def create_gradio_ui():
+    """Create the complete Gradio interface"""
+    
+    with gr.Blocks(title="RAG Agent with Agentic Memory", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🤖 RAG Agent with Agentic Memory
+        
+        Upload documents and chat with an intelligent agent that uses:
+        - 📚 **Local Knowledge Base** (ChromaDB)
+        - 🔍 **Web Search** (Tavily)
+        - 📖 **Wikipedia**
+        - 🎓 **ArXiv** (Academic Papers)
+        """)
+        
+        with gr.Tabs():
+            # Documents Tab
+            with gr.Tab("📄 Documents"):
+                gr.Markdown("### Upload and Manage Documents")
+                gr.Markdown("Upload PDF or Markdown files to add them to the knowledge base.")
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        file_upload = gr.File(
+                            label="Upload Documents",
+                            file_count="multiple",
+                            file_types=[".pdf", ".md"]
+                        )
+                        upload_btn = gr.Button("📤 Add to Knowledge Base", variant="primary")
+                        upload_status = gr.Textbox(label="Upload Status", lines=3)
+                    
+                    with gr.Column(scale=1):
+                        file_list = gr.Textbox(
+                            label="Documents in Knowledge Base",
+                            lines=10,
+                            value=get_file_list()
+                        )
+                        refresh_btn = gr.Button("🔄 Refresh List")
+                        clear_btn = gr.Button("🗑️ Clear All Documents", variant="stop")
+                
+                # Connect document management buttons
+                upload_btn.click(
+                    fn=upload_files,
+                    inputs=[file_upload],
+                    outputs=[upload_status, file_list]
+                )
+                
+                refresh_btn.click(
+                    fn=get_file_list,
+                    outputs=[file_list]
+                )
+                
+                clear_btn.click(
+                    fn=clear_database,
+                    outputs=[upload_status, file_list]
+                )
+            
+            # Chat Tab
+            with gr.Tab("💬 Chat"):
+                gr.Markdown("### Chat with Your Documents")
+                gr.Markdown("Ask questions about your documents or any topic. The agent will search multiple sources.")
+                
+                chatbot = gr.Chatbot(
+                    label="Conversation",
+                    height=500,
+                    show_label=True,
+                    avatar_images=(None, "🤖")
+                )
+                
+                with gr.Row():
+                    msg = gr.Textbox(
+                        label="Your Message",
+                        placeholder="Ask me anything about your documents or general knowledge...",
+                        scale=4
+                    )
+                    submit_btn = gr.Button("Send", variant="primary", scale=1)
+                
+                with gr.Row():
+                    clear_chat_btn = gr.Button("🔄 Reset Conversation")
+                    gr.Markdown("*Note: Resetting clears the conversation history*")
+                
+                # Chat interface
+                chat_interface = gr.ChatInterface(
+                    fn=chat_with_agent,
+                    chatbot=chatbot,
+                    textbox=msg,
+                    submit_btn=submit_btn,
+                    retry_btn=None,
+                    undo_btn=None,
+                    clear_btn=None
+                )
+                
+                clear_chat_btn.click(
+                    fn=reset_conversation,
+                    outputs=[chatbot]
+                )
+        
+        gr.Markdown("""
+        ---
+        ### 🔧 How it works:
+        1. **Upload documents** in the Documents tab
+        2. **Ask questions** in the Chat tab
+        3. The agent will:
+           - Analyze your query
+           - Search relevant sources
+           - Provide comprehensive answers with citations
+        """)
+    
+    return demo
+
+if __name__ == "__main__":
+    demo = create_gradio_ui()
+    demo.launch(share=False, server_name="127.0.0.1", server_port=7860)

core/rag_agent.py

@@ -5,15 +5,18 @@ from agent.tools import *
 from agent.graph import create_agent_graph
 
 class RAGAgent:
-    
-    def __init__(self, collection_name=config.CHILD_COLLECTION):
-        self.collection_name = collection_name
-        self.retriever = intialize_chroma_vectorstore()
+    def __init__(self):
         self.thread_id = str(uuid.uuid4())
         
-        self.llm = ChatGoogleGenerativeAI(model=config.LLM_MODEL, temperature=config.LLM_TEMPERATURE)
-        tools = []
-        self.agent_graph = create_agent_graph(self.llm, tools)
+        self.llm = ChatGoogleGenerativeAI(
+            model=config.LLM_MODEL, 
+            temperature=config.LLM_TEMPERATURE
+        )
+
+        vectordb = intialize_chroma_vectorstore()
+
+        search_tools = [web_search_tavily, arxiv_search, wikipedia_search]
+        self.agent_graph = create_agent_graph(self.llm, vectordb, search_tools)
         
     def get_config(self):
         return {"configurable": {"thread_id": self.thread_id}}

main.py

@@ -1,6 +1,11 @@
-def main():
-    print("Hello from rag-agent!")
-
+from ui.gradio_components import create_gradio_ui
 
 if __name__ == "__main__":
-    main()
+    print("🚀 Launching RAG Agent UI...")
+    demo = create_gradio_ui()
+    demo.launch(
+        share=False,
+        server_name="127.0.0.1",
+        server_port=7860,
+        show_error=True
+    )

ui/gradio_components.py

@@ -1,91 +1,116 @@
 import gradio as gr
-from core.chat_interface import ChatInterface
-from core.document_manager import DocumentManager
-from core.rag_system import RAGSystem
+from core.rag_agent import RAGAgent
 
-def create_gradio_ui():
-    rag_system = RAGSystem()
-    rag_system.initialize()
-    
-    doc_manager = DocumentManager(rag_system)
-    chat_interface = ChatInterface(rag_system)
-    
-    def format_file_list():
-        files = doc_manager.get_markdown_files()
-        if not files:
-            return "📭 No documents available in the knowledge base"
-        return "\n".join([f"{f}" for f in files])
+# Initialize components
+rag_agent = None
+
+def initialize_agent():
+    """Initialize RAG agent lazily"""
+    global rag_agent
+    if rag_agent is None:
+        rag_agent = RAGAgent()
+    return rag_agent
+
+def chat_with_agent(message, history):
+    """Handle chat interactions with the RAG agent"""
+    if not message.strip():
+        return history
     
-    def upload_handler(files, progress=gr.Progress()):
-        if not files:
-            return None, format_file_list()
-            
-        added, skipped = doc_manager.add_documents(
-            files, 
-            progress_callback=lambda p, desc: progress(p, desc=desc)
-        )
+    try:
+        agent = initialize_agent()
         
-        gr.Info(f"✅ Added: {added} | Skipped: {skipped}")
-        return None, format_file_list()
-    
-    def clear_handler():
-        doc_manager.clear_all()
-        gr.Info(f"🗑️ Removed all documents")
-        return format_file_list()
-    
-    def chat_handler(msg, hist):
-        return chat_interface.chat(msg, hist)
+        # Stream the agent's response
+        response_text = ""
+        for event in agent.agent_graph.stream(
+            {"messages": [("user", message)]},
+            agent.get_config(),
+            stream_mode="values"
+        ):
+            if "messages" in event and len(event["messages"]) > 0:
+                last_message = event["messages"][-1]
+                if hasattr(last_message, "content"):
+                    response_text = last_message.content
+        
+        if not response_text:
+            response_text = "I apologize, but I couldn't generate a response. Please try again."
+        
+        return response_text
     
-    def clear_chat_handler():
-        chat_interface.clear_session()
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+def reset_conversation():
+    """Reset the conversation thread"""
+    global rag_agent
+    if rag_agent:
+        rag_agent.reset_thread()
+    return None  # Clear chat history
+
+def create_gradio_ui():
+    """Create the complete Gradio interface"""
     
-    with gr.Blocks(title="Agentic RAG") as demo:
+    with gr.Blocks(title="RAG Agent with Agentic Memory", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🤖 RAG Agent with Agentic Memory
         
-        with gr.Tab("Documents", elem_id="doc-management-tab"):
-            gr.Markdown("## Add New Documents")
-            gr.Markdown("Upload PDF or Markdown files. Duplicates will be automatically skipped.")
-            
-            files_input = gr.File(
-                label="Drop PDF or Markdown files here",
-                file_count="multiple",
-                type="filepath",
-                height=200,
-                show_label=False
-            )
-            
-            add_btn = gr.Button("Add Documents", variant="primary", size="md")
-            
-            gr.Markdown("## Current Documents in the Knowledge Base")
-            file_list = gr.Textbox(
-                value=format_file_list(),
-                interactive=False,
-                lines = 7,
-                max_lines=10,
-                elem_id="file-list-box",
-                show_label=False
-            )
-            
-            with gr.Row():
-                refresh_btn = gr.Button("Refresh", size="md")
-                clear_btn = gr.Button("Clear All", variant="stop", size="md")
-            
-            add_btn.click(
-                upload_handler, 
-                [files_input], 
-                [files_input, file_list], 
-                show_progress="corner"
-            )
-            refresh_btn.click(format_file_list, None, file_list)
-            clear_btn.click(clear_handler, None, file_list)
+        Chat with an intelligent agent that uses:
+        - 📚 **Local Knowledge Base** (ChromaDB)
+        - 🔍 **Web Search** (Tavily)
+        - 📖 **Wikipedia**
+        - 🎓 **ArXiv** (Academic Papers)
+        """)
+        
+        gr.Markdown("### Chat with Your Documents")
+        gr.Markdown("Ask questions about your documents or any topic. The agent will search multiple sources.")
         
-        with gr.Tab("Chat"):
-            chatbot = gr.Chatbot(
-                height=600, 
-                placeholder="Ask me anything about your documents!",
-                show_label=False
+        chatbot = gr.Chatbot(
+            label="Conversation",
+            height=500,
+            show_label=True,
+            avatar_images=(None, "🤖")
+        )
+        
+        with gr.Row():
+            msg = gr.Textbox(
+                label="Your Message",
+                placeholder="Ask me anything about your documents or general knowledge...",
+                scale=4
             )
-            chatbot.clear(clear_chat_handler)
-            
-            gr.ChatInterface(fn=chat_handler, chatbot=chatbot)
+            submit_btn = gr.Button("Send", variant="primary", scale=1)
+        
+        with gr.Row():
+            clear_chat_btn = gr.Button("🔄 Reset Conversation")
+            gr.Markdown("*Note: Resetting clears the conversation history*")
+        
+        # Chat interface
+        chat_interface = gr.ChatInterface(
+            fn=chat_with_agent,
+            chatbot=chatbot,
+            textbox=msg,
+            submit_btn=submit_btn,
+            retry_btn=None,
+            undo_btn=None,
+            clear_btn=None
+        )
+        
+        clear_chat_btn.click(
+            fn=reset_conversation,
+            outputs=[chatbot]
+        )
+        
+        gr.Markdown("""
+        ---
+        ### 🔧 How it works:
+        1. **Ask questions** in the chat
+        2. The agent will:
+           - Analyze your query
+           - Search relevant sources (ChromaDB, Web, Wikipedia, ArXiv)
+           - Provide comprehensive answers with citations
+        3. Use **Reset Conversation** to start fresh
+        """)
     
-    return demo
+    return demo
+
+if __name__ == "__main__":
+    demo = create_gradio_ui()
+    demo.launch(share=False, server_name="127.0.0.1", server_port=7860)