update files

agents/agent_graph.py

@@ -23,7 +23,7 @@ def chat_node(state):
 
 def documentsearch_node(state):
     print("[AgentGraph] Invoking DocumentSearchAgent")
-    agent = DocumentSearchAgent()
+    agent = DocumentSearchAgent(urls=state.get('doc_urls'), pdf_files=state.get('pdf_files'))
     user_input = state['user_input']
     chat_history = state.get('chat_history', [])
     response, context, trace = agent.run(user_input, chat_history)
@@ -80,6 +80,9 @@ def finalize_node(state):
     """
     Finalize response using OrchestratorAgent.system_prompt.
     """
+    # If document search indicated no sources, skip refinement
+    if any(entry.get('step') == 'no_sources' for entry in state.get('trace', [])):
+        return state
     return orchestrator.finalize_response(state)
 
 def input_guard_node(state):
@@ -89,6 +92,7 @@ def input_guard_node(state):
     return orchestrator.validate_input(state)
 
 def orchestrator_node(state):
+    print("[AgentGraph] Invoking OrchestratorAgent")
     user_input = state['user_input']
     chat_history = state.get('chat_history', [])
     agent_name = orchestrator._decide_agent(user_input, chat_history=chat_history)
@@ -148,13 +152,13 @@ def build_agent_graph():
     graph.add_edge('chat', 'finalize')
     graph.add_edge('websearch', 'finalize')
     graph.add_edge('documentsearch', 'evaluate')
-    # Evaluate documentsearch and fallback based on search toggles
+    # Evaluate documentsearch and fallback based on search toggles, but if no_sources, finalize immediately
     graph.add_conditional_edges(
         'evaluate',
         lambda state: (
-            'finalize' if state.get('answered', False)
+            'finalize' if any(step.get('step') == 'no_sources' for step in state.get('trace', []))
+            else 'finalize' if state.get('answered', False)
             else 'websearch' if state.get('enable_websearch', True)
-            # else 'documentsearch' if state.get('enable_docsearch', True)
             else 'chat'
         ),
         {

agents/chat/chat_agent.py

@@ -16,7 +16,7 @@ class ChatAgent(AgentBase):
         prompt = PromptTemplate(
             input_variables=["input"],
             template=(
-                "You are an assistant specialized in Thang Long University information. Use the provided tools when appropriate.\n"
+                "You are Kia, the Know-It-All assistant. Use the provided tools when appropriate.\n"
                 "\nQuestion: {input}\n{agent_scratchpad}"
             )
         )
@@ -106,10 +106,15 @@ class ChatAgent(AgentBase):
             user_input = state['user_input']
             history_str = state.get('history_str', '')
             prompt = (
-                f"You are a helpful assistant. Here is the conversation so far:\n"
-                f"{history_str}"
-                f"User: {user_input}\nBot:"
+                "You are Kia, the Know-It-All assistant.\n"    
+                "Your task is to answer the user's question.\n"
+                "Here is the conversation so far:\n"
+                f"{history_str}\n"
+                "Here is the user's question:\n"
+                f"{user_input}\n"
+                "Your response:\n"
             )
+
             state['prompt'] = prompt
             state.setdefault('trace', []).append({'step': 'build_prompt', 'prompt': prompt})
             return state
@@ -119,12 +124,13 @@ class ChatAgent(AgentBase):
             messages = [HumanMessage(content=state['prompt'])]
             # res = self.llm.chat_model.invoke(messages)
             res = self.agent_executor.invoke({"input": state['prompt']})
+            # print('res', res)
             output = res.get("output")
             actions = res.get("actions", [])
             state['response'] = output
             state['handled'] = True
             state['actions'] = actions
-            state.setdefault('trace', []).append({'step': 'llm_call', 'response': output, 'actions': actions, 'agent': 'chat'})
+            state.setdefault('trace', []).append({'step': 'llm_call', 'response': output, 'actions': actions})
             return state
 
 
@@ -171,14 +177,19 @@ class ChatAgent(AgentBase):
 
     def run(self, user_input: str, chat_history=None):
         # Try the tool-enabled agent executor first
-        try:
-            # Use invoke to capture function calls and args
-            result = self.agent_executor.invoke({"input": user_input})
-            output = result.get("output")
-            actions = result.get("actions", [])
-            return output, actions
-        except Exception:
-            # Fallback to the original graph-based flow
-            state = {'user_input': user_input, 'chat_history': chat_history or [], 'trace': []}
-            result = self.graph.invoke(state)
-            return result['response'], result['trace']
+        # Fallback to the original graph-based flow
+        state = {'user_input': user_input, 'chat_history': chat_history or [], 'trace': []}
+        result = self.graph.invoke(state)
+        return result['response'], result['trace']
+
+        # try:
+        #     # Use invoke to capture function calls and args
+        #     result = self.agent_executor.invoke({"input": user_input})
+        #     output = result.get("output")
+        #     actions = result.get("actions", [])
+        #     return output, actions
+        # except Exception:
+        #     # Fallback to the original graph-based flow
+        #     state = {'user_input': user_input, 'chat_history': chat_history or [], 'trace': []}
+        #     result = self.graph.invoke(state)
+        #     return result['response'], result['trace']

agents/documentsearch/documentsearch_agent.py

@@ -1,4 +1,4 @@
-from langchain_community.document_loaders import WebBaseLoader
+from langchain_community.document_loaders import WebBaseLoader, PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_community.vectorstores import FAISS
@@ -14,26 +14,58 @@ def get_embeddings():
     return embeddings
 
 class DocumentSearchAgent(AgentBase):
-    def __init__(self, urls=None):
-        # Crawl & index website content
-        if urls is None:
-            urls = ["https://en.thanglong.edu.vn/"]
-        loader = WebBaseLoader(urls)
-        documents = loader.load()
+    def __init__(self, urls=None, pdf_files=None):
+        # Prepare raw documents list (only from provided URLs or PDFs)
+        docs_raw = []
+        # Load from URLs
+        if urls:
+            loader = WebBaseLoader(urls)
+            docs_raw.extend(loader.load())
+        # Load from uploaded PDF files
+        if pdf_files:
+            import tempfile, os
+            for uploaded_file in pdf_files:
+                suffix = os.path.splitext(uploaded_file.name)[1] or ".pdf"
+                tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+                tmp.write(uploaded_file.read())
+                tmp.close()
+                loader_pdf = PyPDFLoader(tmp.name)
+                docs_raw.extend(loader_pdf.load())
+        # If no sources provided, skip indexing and prepare guard-only graph
+        if not docs_raw:
+            self.vectorstore = None
+            self.retriever = None
+            self.graph = self.build_graph()
+            return
+        # Split documents into chunks and build index
         splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        docs = splitter.split_documents(documents)
+        docs = splitter.split_documents(docs_raw)
         embeddings = get_embeddings()
         self.vectorstore = FAISS.from_documents(docs, embeddings)
         self.retriever = self.vectorstore.as_retriever(search_kwargs={"k": 5})
         self.graph = self.build_graph()
+        self.llm = LLM()
 
     def build_graph(self):
+        # Guard node: check for document sources
+        def check_sources(state):
+            if self.retriever is None:
+                msg = "😞 Tính năng tìm kiếm tài liệu đã bật, nhưng bạn chưa tải lên tài liệu hoặc nhập URL. Vui lòng bổ sung để tiếp tục."
+                state['response'] = msg
+                state['handled'] = True
+                state.setdefault('trace', []).append({'step': 'no_sources', 'response': msg})
+            return state
         def rephrase(state):
-            # (Optional) Could implement query rephrasing using LLM here
-            state.setdefault('trace', []).append({'step': 'rephrase', 'query': state['user_input']})
+            history_str = ""
+            chat_history = state.get('chat_history', [])
+            for turn in chat_history:
+                history_str += f"User: {turn['user']}\nBot: {turn['bot']}\n"
+            rephrased_query = self.llm.generate(f"User: {state['user_input']}\nBot: Rephrase the user's question so it is clear and complete for a document search. Only output the rephrased question.\n{history_str}").strip()
+            state['rephrased_query'] = rephrased_query
+            state.setdefault('trace', []).append({'step': 'rephrase', 'query': state['user_input'], 'rephrased_query': rephrased_query})
             return state
         def retrieve(state):
-            query = state['user_input']
+            query = state['rephrased_query']
             results = self.retriever.invoke(query)
             # print("len(results)", len(results))
             # Store raw results instead of joining them
@@ -45,6 +77,7 @@ class DocumentSearchAgent(AgentBase):
             context = state.get('context', [])
             # print("len(context)", len(context))
             user_input = state['user_input']
+            rephrased_query = state['rephrased_query']
             llm = LLM()
 
             # Format context as markdown table, one row per document
@@ -59,25 +92,36 @@ class DocumentSearchAgent(AgentBase):
                 context_table = "| # | Content |\n|---|---------|\n" + "\n".join(table_rows)
 
             prompt = (
-                f"Based on the following excerpts from https://en.thanglong.edu.vn/, answer the user's question as accurately as possible.\n\n"
+                "You are Kia, the Know-It-All assistant. Based on the following provided document excerpts, answer the user's question as accurately and concisely as possible.\n\n"
                 f"Excerpts (in markdown table format):\n{context_table}\n\n"
-                f"User: {user_input}\nBot:"
+                f"User: {rephrased_query}\nBot:"
             )
             response = llm.generate(prompt)
             state['response'] = response
             state.setdefault('trace', []).append({'step': 'synthesize', 'prompt': prompt, 'response': response})
             return state
         graph = StateGraph(dict)
+        graph.add_node('check_sources', check_sources)
         graph.add_node('rephrase', rephrase)
         graph.add_node('retrieve', retrieve)
         graph.add_node('synthesize', synthesize)
+        # If no sources, exit immediately
+        graph.add_conditional_edges(
+            'check_sources',
+            lambda state: state.get('handled', False),
+            {True: END, False: 'rephrase'}
+        )
         graph.add_edge('rephrase', 'retrieve')
         graph.add_edge('retrieve', 'synthesize')
         graph.add_edge('synthesize', END)
-        graph.set_entry_point('rephrase')
+        graph.set_entry_point('check_sources')
         return graph.compile()
 
     def run(self, query: str, chat_history=None):
         state = {'user_input': query, 'chat_history': chat_history or [], 'trace': []}
         result = self.graph.invoke(state)
-        return result['response'], result['context'], result['trace']
+        # Safely extract response, context, and trace
+        response = result.get('response', '')
+        context = result.get('context', [])
+        trace = result.get('trace', [])
+        return response, context, trace

agents/orchestrator/orchestrator_agent.py

@@ -11,11 +11,9 @@ class OrchestratorAgent:
         self.llm = LLM()
         # System prompt for all agents
         self.system_prompt = (
-            "You are an AI assistant specialized in Thang Long University information. "
-            "Provide accurate, concise, and clear responses in a professional style. "
+            "You are Kia, the Know-It-All assistant. Provide accurate, concise, and clear responses in a professional style. "
             "Respond in the user's language. "
-            "Do not provide personal or sensitive data, refuse harmful or inappropriate requests, "
-            "avoid off-topic responses (politely decline queries outside Thang Long University domain). "
+            "Do not provide personal or sensitive data, and refuse harmful or inappropriate requests."
         )
 
     def _decide_agent(self, user_input: str, chat_history=None) -> str:
@@ -80,15 +78,13 @@ class OrchestratorAgent:
             context_table = "| Role | Message |\n|------|---------|\n" + "\n".join(table_rows)
 
         prompt = (
-            "As an AI assistant, analyze the user's query and conversation context to determine if it requires searching through Thang Long University's website documents.\n\n"
+            "As an AI assistant, analyze the user's query and conversation context to determine if it requires searching through the provided documents (uploaded PDFs or URLs).\n\n"
             "Context:\n"
             f"Recent conversation:\n{context_table}\n\n"
             f"Current query: {user_input}\n\n"
             "Consider these factors:\n"
-            "1. Is the query specifically about Thang Long University?\n"
-            "2. Does it require accessing structured document content?\n"
-            "3. Is the information likely to be found in university documentation?\n"
-            "4. Would document search be more reliable than web search for this query?\n\n"
+            "1. Does the query require locating information within the provided documents?\n"
+            "2. Is document search more reliable than web search for this query?\n\n"
             "Output ONLY 'documentsearch' if document search is most appropriate, or 'other' if not."
         )
 
@@ -117,25 +113,13 @@ class OrchestratorAgent:
         if table_rows:
             history_table = "| Role | Message |\n|------|---------|\n" + "\n".join(table_rows)
 
-        # Old routing prompt (commented out):
-        # prompt = (
-        #     "You are an orchestrator for a multi-agent assistant. "
-        #     "Decide which agent should handle the user's latest message: "
-        #     "output ONLY 'chat' (for general conversation, Q&A, reasoning, etc.) "
-        #     "or 'websearch' (if the user is asking for real-time, factual, or web-based information.)\n\n"
-        #     f"Conversation so far:\n{history_table}\n\n"
-        #     f"Current query: {user_input}\n\n"
-        #     "Which agent should handle this? (chat/websearch):"
-        # )
-        # decision = self.llm.generate(prompt).strip().lower()
-        # return "websearch" if "websearch" in decision else "chat"
-        # New custom prompt for Thang Long University:
+        # New custom prompt for general assistant:
         prompt = (
-            "You are an orchestrator specialized in Thang Long University information. "
+            "You are an orchestrator for a general AI assistant. "
             "Route the user's query to the most suitable agent: "
-            "'chat' for follow-up or general queries, "
+            "'chat' for general conversation or reasoning, "
             "'websearch' for real-time or external data, "
-            "'documentsearch' for internal university documents. "
+            "'documentsearch' for searching within provided documents (PDFs/URLs). "
             f"Conversation history:\n{history_table}\n\n"
             f"User query: {user_input}\n\n"
             "Respond with ONLY the agent name (chat/websearch/documentsearch)."
@@ -162,12 +146,12 @@ class OrchestratorAgent:
             "websearch_results": []
         }
 
-    def finalize_response(self, result_state: dict) -> dict:
+    def finalize_response(self, state: dict) -> dict:
         """
         Refine the routed agent's output using the system prompt.
         """
-        raw = result_state.get('response', '')
-        question = result_state.get('user_input', '')
+        raw = state.get('response', '')
+        question = state.get('user_input', '')
         # Construct detailed refinement prompt
         prompt = (
             f"{self.system_prompt}\n\n"
@@ -176,18 +160,30 @@ class OrchestratorAgent:
             f"{raw}\n\n"
             "Please refine this answer: polite, professional, and in the user's language; focus solely on answering the question without extraneous remarks. "
             "Return only the final answer text."
+            "Chỉ trong trường hợp draft answer chứa nhiều mâu thuẫn, không trả lời đầy đủ, hoặc không chính xác, hãy tinh chỉnh. "
+            "Nếu không, hãy giữ nguyên."
         )
-        refined = self.llm.generate(prompt).strip()
-        result_state['response'] = refined
-        # Label finalized output under orchestrator
-        result_state['agent'] = 'orchestrator'
-        result_state.setdefault('trace', []).append({
+
+        # refined = self.llm.generate(prompt).strip()
+        # result_state['response'] = refined
+        # # Label finalized output under orchestrator
+        # result_state['agent'] = 'orchestrator'
+        # result_state.setdefault('trace', []).append({
+        #     'step': 'finalize',
+        #     'prompt': prompt,
+        #     'response': refined,
+        #     'agent': 'orchestrator'
+        # })
+
+        state['response'] =  raw
+        state['agent'] = 'orchestrator'
+        state.setdefault('trace', []).append({
             'step': 'finalize',
             'prompt': prompt,
-            'response': refined,
+            'response': raw,
             'agent': 'orchestrator'
         })
-        return result_state
+        return state
 
     def validate_input(self, state: dict) -> dict:
         """
@@ -199,7 +195,7 @@ class OrchestratorAgent:
         # Empty input
         if not user_input:
             state['guard_failed'] = True
-            state['response'] = "😞 Sorry, I didn't catch that. Please ask a question about Thang Long University."
+            state['response'] = "😞 Sorry, I didn't catch that. Please ask a question."
             state['agent'] = 'orchestrator'
             # Record validation failure
             state.setdefault('trace', []).append({'step': 'validate', 'reason': 'empty input', 'agent': 'orchestrator'})
@@ -212,15 +208,15 @@ class OrchestratorAgent:
             history_str += f"User: {turn['user']}\nBot: {turn['bot']}\n"
         classification_prompt = (
             f"Given the following conversation history:\n{history_str}"
-            "Classify the following query as ON_TOPIC or OFF_TOPIC. Domain: Thang Long University information.\n"
+            "Classify the following query as ON_TOPIC or OFF_TOPIC with respect to a general conversational assistant.\n"
             f"Query: \"{user_input}\"\n"
-            "General greetings are not considered off-topic."
+            "General greetings are not considered off-topic.\n"
             "Respond with exactly ON_TOPIC or OFF_TOPIC."
         )
         classification = self.llm.generate(classification_prompt).strip().upper()
         if classification != "ON_TOPIC":
             state['guard_failed'] = False
-            state['response'] = "😞 Sorry, I can only answer questions about Thang Long University."
+            state['response'] = "😞 Sorry, I can only answer on-topic questions."
             state['agent'] = 'orchestrator'
             # Record off-topic validation
             state.setdefault('trace', []).append({'step': 'validate', 'reason': 'off topic_lm', 'agent': 'orchestrator'})

agents/websearch/websearch_agent.py

@@ -17,7 +17,7 @@ class WebSearchAgent(AgentBase):
                 history_str += f"User: {turn['user']}\nBot: {turn['bot']}\n"
             state['history_str'] = history_str
             rephrase_prompt = (
-                f"You are a helpful assistant specialized in Thang Long University information.\n"
+                f"You are a helpful assistant.\n"
                 f"Given the following conversation history:\n"
                 f"{history_str}"
                 f"User: {state['user_input']}\n"

app.py

@@ -13,13 +13,13 @@ import time
 def generate_response_message(response):
     full_response = ""
     response_words = response.split()
-    with st.chat_message("Kia", avatar="🤖"):
-        message_placeholder = st.empty()
-        for word in response_words:
-            full_response += word + " "
-            message_placeholder.markdown(full_response + "▌")
-            time.sleep(0.05)
-        message_placeholder.markdown(full_response)
+    # with st.chat_message("Kia", avatar="🤖"):
+    message_placeholder = st.empty()
+    for word in response_words:
+        full_response += word + " "
+        message_placeholder.markdown(full_response + "▌")
+        time.sleep(0.05)
+    message_placeholder.markdown(full_response)
     return full_response
 
 # Function to generate initial message
@@ -59,8 +59,13 @@ with st.sidebar:
     if 'enable_websearch' not in st.session_state:
         st.session_state['enable_websearch'] = True
     with st.expander("Agent Settings", expanded=True):
-        
         st.checkbox("Enable Document Search", value=st.session_state['enable_docsearch'], key='enable_docsearch')
+        if st.session_state['enable_docsearch']:
+            st.write("Upload PDF(s) or enter URL(s) to index for document search")
+            st.file_uploader("Upload PDF document(s)", type=["pdf"], accept_multiple_files=True, key="pdf_files")
+            st.text_input("Enter URL(s), comma-separated", key="doc_urls")
+            if st.button("Rebuild Index"):
+                st.experimental_rerun()
         st.checkbox("Enable Web Search", value=st.session_state['enable_websearch'], key='enable_websearch')
     
     # Chat Controls
@@ -73,7 +78,7 @@ with st.sidebar:
     # Settings Section
     with st.expander("🛠️ Display Settings", expanded=True):
         DEV_MODE = st.checkbox("Enable Dev Mode", value=False)
-        STREAMING_ENABLED = st.checkbox("Enable response streaming", value=True)
+        STREAMING_ENABLED = st.checkbox("Enable response streaming (just the effect 😂)", value=True)
         # # Theme selector
         # if 'theme' not in st.session_state:
         #     st.session_state.theme = "Light"
@@ -177,11 +182,56 @@ for i, turn in enumerate(st.session_state['chat_history']):
         continue
     # Assistant response rendering (streaming optional)
     if i == len(st.session_state['chat_history']) - 1 and STREAMING_ENABLED and st.session_state['should_stream']:
-        generate_response_message(turn['bot'])
+        with st.chat_message("Kia", avatar="🤖"):
+            generate_response_message(turn['bot'])
+
+            # Display sources for websearch in streaming mode
+            if turn.get('trace') is not None:
+                for step in turn['trace']:
+                    if step.get('agent') == 'websearch' and step.get('step') == 'search':
+                        raw_results = step.get('raw_results')
+                        if raw_results:
+                            st.markdown("**Sources:**")
+                        if isinstance(raw_results, dict) and raw_results.get('organic'):
+                            for idx, item in enumerate(raw_results['organic'], 1):
+                                link = item.get('link')
+                                title = item.get('title', '')
+                                st.markdown(f"{idx}. [{title}]({link})")
+                        elif isinstance(raw_results, list):
+                            for idx, item in enumerate(raw_results, 1):
+                                link = item.get('link') if isinstance(item, dict) else item
+                                title = item.get('title', '') if isinstance(item, dict) else ''
+                                st.markdown(f"{idx}. [{title}]({link})")
+                        break
         st.session_state['should_stream'] = False
     else:
         with st.chat_message("Kia", avatar="🤖"):
             st.markdown(turn['bot'])
+            # st.markdown(turn['trace'])
+
+            if turn['trace'] is not None:
+                for step in turn['trace']:
+                    if step.get('agent') == 'websearch' and step.get('step') == 'search':
+                        # Extract raw_results from trace
+                        raw_results = step.get('raw_results')
+                        # Display citation sources for websearch agent
+                        if raw_results:
+                            st.markdown("**Sources:**")
+                        # Prefer organic list from raw_results
+                        if isinstance(raw_results, dict) and raw_results.get('organic'):
+                            for idx, item in enumerate(raw_results['organic'], 1):
+                                link = item.get('link')
+                                title = item.get('title', '')
+                                st.markdown(f"{idx}. [{title}]({link})")
+                        # Fallback if raw_results itself is list
+                        elif isinstance(raw_results, list):
+                            for idx, item in enumerate(raw_results, 1):
+                                link = item.get('link') if isinstance(item, dict) else item
+                                title = item.get('title', '') if isinstance(item, dict) else ''
+                                st.markdown(f"{idx}. [{title}]({link})")
+                        break
+
+    
     # DEV_MODE trace expander (all completed turns)
     if DEV_MODE and turn.get('trace'):
         trace = turn.get('trace', [])
@@ -234,11 +284,17 @@ if st.session_state['waiting_for_response']:
     user_input = st.session_state['chat_history'][-1]['user']
     agent_graph = build_agent_graph()
     chat_history = st.session_state['chat_history'][:-1]  # Exclude the waiting message
+    doc_urls = None
+    if st.session_state.get('doc_urls'):
+        doc_urls = [u.strip() for u in st.session_state['doc_urls'].split(",") if u.strip()]
+    pdf_files = st.session_state.get('pdf_files', [])
     state = {
         "user_input": user_input,
         "chat_history": chat_history,
         "enable_docsearch": st.session_state['enable_docsearch'],
         "enable_websearch": st.session_state['enable_websearch'],
+        "doc_urls": doc_urls,
+        "pdf_files": pdf_files,
     }
 
     # --- Invoke agent graph (may take time) with error handling ---

main.py

@@ -5,7 +5,7 @@ def main():
     # Visualize the agent graph every time it is built
     visualize_agent_graph(agent_graph, as_image=True, save_to_file=True)
     chat_history = []
-    print("Multi-Agent Chatbot (type 'exit' to quit)")
+    print("Kia - Your Know-It-All Assistant (type 'exit' to quit)")
     while True:
         user_input = input("You: ")
         if user_input.lower() in ["exit", "quit"]: