add summary table

agent.py

@@ -19,7 +19,7 @@ from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
 from langchain.chat_models import init_chat_model
 from langchain.agents import initialize_agent, AgentType
 from langchain_community.retrievers import BM25Retriever
-from langchain.schema import BaseMessage, SystemMessage, HumanMessage
+from langchain.schema import BaseMessage, SystemMessage, HumanMessage, AIMessage
 from langgraph.graph.message import add_messages
 from langgraph.graph import START, END, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
@@ -284,6 +284,8 @@ def extract_table(file_path: str, query: str = "") -> str:
         df = pd.read_csv(file_path)
     elif ext in [".xlsx", ".xls"]:
         df = pd.read_excel(file_path)
+        text_content = df.to_string()
+        loaded_docs = [Document(page_content=text_content)]
     else:
         return "Unsupported file type."
     # Simple filter: return all if no query, else filter columns containing query
@@ -292,12 +294,23 @@ def extract_table(file_path: str, query: str = "") -> str:
         df = df[mask]
     return df.head(10).to_csv(index=False)
 
+@tool
+def summarize(text: str, llm=None) -> str:
+    """Summarize a long text chunk."""
+    if llm is None:
+        return "No LLM provided for summarization."
+    return llm.invoke([
+        SystemMessage(content="Summarize the following:"),
+        HumanMessage(content=text)
+    ]).content
+
 # Update tools list
 tools: List[StructuredTool] = [
     calculate, tavily_search, wikipedia_search, image_recognition,
     read_pdf, read_csv, read_spreadsheet, transcribe_audio,
     youtube_transcript_tool, youtube_transcript_api, read_jsonl,
-    python_interpreter, download_file, extract_table   # Add tavily_search here
+    python_interpreter, download_file, extract_table,
+    # Wrap summarize to inject self.llm at runtime
 ]
 
 class AgentState(TypedDict):
@@ -317,9 +330,11 @@ class MyAgent:
             self.llm = init_chat_model(
                 model_name,
                 temperature=temperature
-                )
+            )
             # Base tools
-            self.tools = tools
+            self.tools = tools + [
+                StructuredTool.from_function(lambda text: summarize(text, llm=self.llm), name="summarize", description="Summarize a long text chunk.")
+            ]
             # RAG components
             self.docs: List[Any] = []
             self.retriever: Optional[BM25Retriever] = None
@@ -334,51 +349,49 @@ class MyAgent:
         """
         for path in file_paths:
             ext = Path(path).suffix.lower()
+            loaded_docs = []
             try:
                 if ext == ".csv":
                     loader = CSVLoader(path)
-                    self.docs.extend(loader.load())
+                    loaded_docs = loader.load()
                 elif ext == ".pdf":
                     loader = PyPDFLoader(path)
-                    self.docs.extend(loader.load())
+                    loaded_docs = loader.load()
                 elif ext in [".xlsx", ".xls"]:
-                    # Handle spreadsheets
                     import pandas as pd
                     df = pd.read_excel(path)
                     text_content = df.to_string()
-                    self.docs.append(Document(page_content=text_content))
+                    loaded_docs = [Document(page_content=text_content)]
                 elif ext == ".jsonl":
-                    # Handle JSONL files
                     with open(path, 'r', encoding='utf-8') as file:
                         content = [json.loads(line) for line in file]
                         text_content = json.dumps(content, indent=2)
-                        self.docs.append(Document(page_content=text_content))
+                        loaded_docs = [Document(page_content=text_content)]
                 elif ext in [".png", ".jpg", ".jpeg"]:
-                    # Handle images
                     text = pytesseract.image_to_string(Image.open(path))
                     if text.strip():
-                        self.docs.append(Document(page_content=text))
+                        loaded_docs = [Document(page_content=text)]
                 elif ext in [".mp3", ".wav"]:
                     loader = AssemblyAIAudioTranscriptLoader(file_path=path)
-                    self.docs.extend(loader.load())
+                    loaded_docs = loader.load()
                 elif "youtube" in path:
                     loader = YoutubeLoader.from_youtube_url(path)
-                    self.docs.extend(loader.load())
+                    loaded_docs = loader.load()
                 else:
                     print(f"Unsupported file type: {ext}")
                     continue
             except Exception as e:
                 print(f"Error loading {path}: {e}")
                 continue
-        # After loading each doc:
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
-        for doc in loaded_docs:
-            chunks = text_splitter.split_text(doc.page_content)
-            for i, chunk in enumerate(chunks):
-                self.docs.append(Document(
-                    page_content=chunk,
-                    metadata={**doc.metadata, "chunk": i, "source": path}
-                ))
+            # Chunk every loaded doc
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
+            for doc in loaded_docs:
+                chunks = text_splitter.split_text(doc.page_content)
+                for i, chunk in enumerate(chunks):
+                    self.docs.append(Document(
+                        page_content=chunk,
+                        metadata={**getattr(doc, 'metadata', {}), "chunk": i, "source": path}
+                    ))
 
     def build_retriever(self):
         """
@@ -414,63 +427,55 @@ class MyAgent:
         file_paths: Optional[List[str]] = None
     ) -> str:
         try:
-            # Prepare state graph
-            state: Dict[str, Any] = {"messages": [], "input_file": None}
-
-            # Use structured tool attributes
+            state: Dict[str, Any] = {"messages": [], "input_file": None, "rag_used": False}
             tool_desc = "\n".join(f"{t.name}: {t.description}" for t in self.tools)
-            
-            # Enhanced system prompt with RAG guidance
             rag_prompt = """
             If the question seems to be about any loaded documents, ALWAYS:
             1. Use the rag_search tool first to find relevant information
             2. Base your answer on the retrieved content
             3. If no relevant content is found, say so
             """
-            
             sys_msg = SystemMessage(content=f"{SYSTEM_PROMPT}\n\n{rag_prompt if file_paths else ''}\n\nTools:\n{tool_desc}")
-            state["messages"].append(sys_msg)
-
-            # Optionally load RAG docs
-            if file_paths:
-                self.add_files(file_paths)
-                self.build_retriever()
-
-            # Add user question
+            state["messages"] = [sys_msg]
+            if file_paths and all(isinstance(p, str) for p in file_paths):
+                try:
+                    self.add_files(file_paths)
+                    self.build_retriever()
+                except Exception as file_err:
+                    print(f"Warning: Error loading files: {file_err}")
             state["messages"].append(HumanMessage(content=question))
             if file_paths:
                 state["input_file"] = file_paths
-
-            # Build graph with proper conditional edge to prevent loops
             builder = StateGraph(dict)
             builder.add_node("assistant", self._assistant_node)
-            builder.add_node("tools", ToolNode(self.tools))
+            # Add the tools node BEFORE adding edges
+            def tool_node_with_rag_flag(state):
+                state = ToolNode(self.tools).invoke(state)
+                if state.get("input_file") and not state.get("rag_used", False):
+                    state["rag_used"] = True
+                return state
+            builder.add_node("tools", tool_node_with_rag_flag)
             builder.add_edge(START, "assistant")
-            
-            # Fix conditional edges with better check
-            def _should_use_tools(state):
-                # If there are loaded docs, always use rag_search first
-                if state.get("input_file"):
+            # Graph flow: force rag_search if files loaded and not yet used, then use tools_condition
+            def route(state):
+                # If files loaded and rag not used, force rag_search
+                if state.get("input_file") and not state.get("rag_used", False):
                     return "tools"
-                # Otherwise, let the assistant try to answer
-                return "assistant"
-
-            builder.add_conditional_edges(
-                "assistant",
-                _should_use_tools,
-                {"tools": "tools", "assistant": END}
-            )
+                    
+                last_msg = state["messages"][-1] if state.get("messages") else None
+                # Only route to tools if the last message is an AIMessage and has tool_calls
+                if last_msg and isinstance(last_msg, AIMessage):
+                    if getattr(last_msg, "tool_calls", None):
+                        return "tools"
+                    if getattr(last_msg, "additional_kwargs", {}).get("tool_calls"):
+                        return "tools"
+                return END
+            builder.add_conditional_edges("assistant", route, {"tools": "tools", END: END})
             builder.add_edge("tools", "assistant")
-            
-            # Add recursion_limit to prevent infinite loops
+            # Instead of builder.update_node, define a custom tool node with rag flag logic
             graph = builder.compile()
-            
-            # Use invoke() with higher recursion limit
-            out = graph.invoke(state, {"recursion_limit": 10})  # Lower limit
-            last_message = out["messages"][-1].content
-            
-            # Extract only the FINAL ANSWER part
-            import re
+            out = graph.invoke(state, {"recursion_limit": 10})
+            last_message = out["messages"][-1].content if out.get("messages") else ""
             match = re.search(r"FINAL ANSWER[:\s]*([^\n]*)", last_message, re.IGNORECASE)
             if match:
                 return match.group(1).strip()