CoolFace
Apppublic

DevForML/Multi_Agent_System

sourceHugging Facemitupdated 1y agoView on Hugging Face
0likes
agent.py711 linesDownload Raw Back to root
1#─── Basic imports ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
2import os
3import math
4import sqlite3
5import fitz  # PyMuPDF for PDF parsing
6from flask_socketio import SocketIO
7
8# ─── Langchain Frameworks ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
9from langchain.tools import Tool
10from langchain.chat_models import ChatOpenAI
11from langchain_groq import ChatGroq
12from langchain_mistralai import ChatMistralAI
13from langchain.agents import initialize_agent, AgentType
14from langchain.schema import Document
15from langchain.chains import RetrievalQA
16from langchain.embeddings import OpenAIEmbeddings
17from langchain_community.embeddings import HuggingFaceEmbeddings
18from langchain.vectorstores import FAISS
19from langchain.text_splitter import RecursiveCharacterTextSplitter
20from langchain.prompts import PromptTemplate
21from langchain_community.document_loaders import TextLoader, PyMuPDFLoader
22# taking global variables from the app.py file
23#from app import DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH
24
25# ─── File paths ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
26import config
27# Ensure this is at the very top
28
29# ─── SQL Agent ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
30from langchain_community.utilities import SQLDatabase
31from langchain_community.agent_toolkits import SQLDatabaseToolkit
32from langchain.chat_models import ChatOpenAI
33from langgraph.prebuilt import create_react_agent
34from langchain.agents import create_sql_agent
35
36# ─── Memory ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
37from langchain.memory import ConversationBufferMemory
38from langchain.agents import initialize_agent, AgentType
39from langchain.tools import Tool
40from typing import List, Callable
41from langchain.memory import ConversationBufferMemory
42from langchain.schema import BaseMemory, AIMessage, HumanMessage, SystemMessage
43from langchain.llms.base import LLM
44from langchain.memory.chat_memory import BaseChatMemory
45from pydantic import PrivateAttr
46from langchain_core.messages import get_buffer_string
47
48# 1) Create your memory object
49from typing import List
50from langchain.memory import ConversationBufferMemory
51from langchain.schema import AIMessage, HumanMessage, SystemMessage
52from langchain.llms.base import LLM
53from langchain.memory.chat_memory import BaseChatMemory
54from pydantic import PrivateAttr
55
56class AutoSummaryMemory(ConversationBufferMemory):
57    _llm: LLM = PrivateAttr()
58    _max_entries: int = PrivateAttr()
59    _reduce_to: int = PrivateAttr()
60    _summary_system_prompt: str = PrivateAttr()
61
62    def __init__(
63        self,
64        llm: LLM,
65        memory_key: str = "chat_history",
66        return_messages: bool = True,
67        max_entries: int = 20,
68        reduce_to: int = 5,
69        summary_system_prompt: str = (
70            "Summarize the following conversation so far in a concise paragraph. "
71            "Keep important facts and questions."
72        )
73    ):
74        super().__init__(memory_key=memory_key, return_messages=return_messages)
75        self._llm = llm  # PrivateAttr
76        self._max_entries = max_entries  # PrivateAttr
77        self._reduce_to = reduce_to  # PrivateAttr
78        self._summary_system_prompt = summary_system_prompt  # PrivateAttr
79
80    def add_memory(self, inputs: dict, outputs: dict) -> None:
81        # Add the new turn as normal
82        super().add_memory(inputs=inputs, outputs=outputs)
83        
84        # Check if memory length exceeded
85        msgs = self.chat_memory.messages
86        if len(msgs) >= self._max_entries:
87            full_text = "\n".join([f"{m.type}: {m.content}" for m in msgs])
88            summary = self._llm.predict(f"{self._summary_system_prompt}\n\n{full_text}")
89
90            recent = msgs[-self._reduce_to:]
91            self.chat_memory.messages = [
92                SystemMessage(content="Conversation summary: " + summary),
93                *recent
94            ]
95            
96            
97# ─── Image Processing ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
98 
99from PIL import Image
100import pytesseract
101from transformers import pipeline
102from groq import Groq
103import config
104import requests
105from io import BytesIO
106from PIL import Image
107from transformers import pipeline, TrOCRProcessor, VisionEncoderDecoderModel
108from PIL import Image
109import requests
110from io import BytesIO
111import base64
112from PIL import UnidentifiedImageError
113
114# ─── Browser var ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
115from typing import List, Dict
116import json
117from io import BytesIO
118from langchain.tools import tool  # or langchain_core.tools
119from playwright.sync_api import sync_playwright
120from duckduckgo_search import DDGS
121from bs4 import BeautifulSoup
122import requests
123
124
125
126from playwright.sync_api import sync_playwright
127# Attempt to import Playwright for dynamic page rendering
128try:
129    from playwright.sync_api import sync_playwright
130    _playwright_available = True
131except ImportError:
132    _playwright_available = False
133
134# Define forbidden keywords for basic NSFW filtering
135_forbidden = ["porn", "sex", "xxx", "nude", "erotic"]
136
137
138# ─── LLM Setup ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
139
140
141# Load OpenAI API key from environment (required for LLM and embeddings)
142import os
143
144# API Keys from .env file
145os.environ.setdefault("OPENAI_API_KEY", "<YOUR_OPENAI_KEY>")  # Set your own key or env var
146os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY", "default_key_or_placeholder")
147os.environ["MISTRAL_API_KEY"] = os.getenv("MISTRAL_API_KEY", "default_key_or_placeholder")
148
149# Tavily API Key
150TAVILY_API_KEY = os.getenv("TAVILY_API_KEY", "default_key_or_placeholder")
151_forbidden = ["nsfw", "porn", "sex", "explicit"]
152_playwright_available = True  # set False to disable Playwright
153
154# Globals for RAG system
155vector_store = None
156rag_chain = None
157DB_PATH = None  # will be set when a .db is uploaded
158DOC_PATH = None  # will be set when a document is uploaded
159IMG_PATH = None  # will be set when an image is uploaded
160OTH_PATH = None  # will be set when an other file is uploaded
161
162
163# ─── LLMS ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
164#llm = ChatOpenAI(model_name="gpt-3.5-turbo", streaming=True, temperature=0)
165llm = ChatGroq(model="meta-llama/llama-4-maverick-17b-128e-instruct", streaming=True, temperature=0)
166#llm = ChatMistralAI(model="mistral-large-latest", streaming=True, temperature=0)
167
168
169# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
170# ─────────────────────────────────────────────── Tool for browsing ────────────────────────────────────────────────────────────────────────
171# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
172
173def tavily_search(query: str, top_k: int = 3) -> List[Dict]:
174    """Call Tavily API and return a list of result dicts."""
175    if not TAVILY_API_KEY:
176        print("[Tavily] No API key set. Skipping Tavily search.")
177        return []
178    url = "https://api.tavily.com/search"
179    headers = {
180        "Authorization": f"Bearer {TAVILY_API_KEY}",
181        "Content-Type": "application/json",
182    }
183    payload = {"query": query, "num_results": top_k}
184    try:
185        resp = requests.post(url, headers=headers, json=payload, timeout=10)
186        resp.raise_for_status()
187        data = resp.json()
188        results = []
189        for item in data.get("results", []):
190            results.append({
191                "title": item.get("title", ""),
192                "url": item.get("url", ""),
193                "snippet": item.get("content", "")[:200],
194                "source": "Tavily"
195            })
196        return results
197    except (requests.exceptions.RequestException, ValueError) as e:
198        print(f"[Tavily] search failed: {e}")
199        return []
200
201def duckduckgo_search(query: str, top_k: int = 3) -> List[Dict]:
202    """Query DuckDuckGo and return up to top_k raw SERP hits."""
203    try:
204        results = []
205        with DDGS() as ddgs:
206            for hit in ddgs.text(query, safesearch="On", max_results=top_k):
207                results.append({
208                    "title": hit.get("title", ""),
209                    "url": hit.get("href") or hit.get("url", ""),
210                    "snippet": hit.get("body", ""),
211                    "source": "DuckDuckGo"
212                })
213                if len(results) >= top_k:
214                    break
215        return results
216    except Exception as e:
217        print(f"[DuckDuckGo] search failed: {e}")
218        return []
219
220def hybrid_web_search(query: str, top_k: int = 3) -> str:
221    """
222    Returns a JSON string with combined Tavily + DuckDuckGo results.
223    Always returns non-empty JSON with at least a placeholder result.
224    """
225    tavily = tavily_search(query, top_k)
226    ddg = duckduckgo_search(query, top_k)
227    combined = tavily + ddg
228
229    # Always return at least a message to avoid agent crashes
230    if not combined:
231        combined = [{
232            "title": "No results found",
233            "url": "",
234            "snippet": f"Could not find suitable web results for '{query}'.",
235            "source": "None"
236        }]
237    output = {"query": query, "results": combined}
238    return json.dumps(output, ensure_ascii=False, indent=2)
239
240def web_search(query: str, top_k: int = 3) -> str:
241    """
242    Full hybrid search with Playwright/BeautifulSoup scraping + Tavily/DuckDuckGo.
243    Always returns valid JSON output.
244    """
245    results: List[Dict] = []
246
247    # Step 1: DuckDuckGo + scraping
248    try:
249        with DDGS() as ddgs:
250            hits = ddgs.text(query, safesearch="On", max_results=top_k)
251    except Exception as e:
252        print(f"[web_search] DuckDuckGo lookup failed: {e}")
253        hits = []
254
255    for hit in hits:
256        url = hit.get("href") or hit.get("url")
257        if not url:
258            continue
259
260        try:
261            with sync_playwright() as pw:
262                browser = pw.chromium.launch(headless=True)
263                page = browser.new_page()
264                page.goto(url, wait_until="domcontentloaded", timeout=15000)
265                html = page.content()
266                browser.close()
267            soup = BeautifulSoup(html, "html.parser")
268            text = soup.get_text(separator=" ", strip=True)
269        except Exception as e:
270            print(f"[web_search] scraping failed for {url}: {e}")
271            continue
272
273        if any(f in text.lower() for f in _forbidden):
274            continue
275
276        excerpt = " ".join(text.split()[:200])
277        results.append({
278            "title": hit.get("title", ""),
279            "url": url,
280            "snippet": hit.get("body", ""),
281            "content": excerpt
282        })
283
284    # Step 2: Parse hybrid Tavily + DDG JSON into list
285    try:
286        raw = hybrid_web_search(query, top_k)
287        parsed = json.loads(raw)
288        other = parsed.get("results", [])
289    except Exception as e:
290        print(f"[web_search] parsing hybrid results failed: {e}")
291        other = []
292
293    # Step 3: Combine and return
294    combined = results + other
295    if not combined:
296        combined = [{
297            "title": "No results found",
298            "url": "",
299            "snippet": f"Could not find suitable content for '{query}'.",
300            "source": "None"
301        }]
302
303    output = {
304        "query": query,
305        "sources_count": len(combined),
306        "results": combined,
307        "sources": list({item.get("url", "") for item in combined if item.get("url")})
308    }
309    return json.dumps(output, ensure_ascii=False, indent=2)
310
311# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
312# ─────────────────────────────────────────────── Tool for calculation ─────────────────────────────────────────────────────────────────────
313# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
314
315def calculate(expr: str) -> str:
316    """
317    Evaluates a mathematical expression safely.
318    Uses Python's numexpr for security and speed:contentReference[oaicite:21]{index=21}.
319    """
320    try:
321        # Allow math constants
322        local_dict = {"pi": math.pi, "e": math.e}
323        # Evaluate expression using numexpr for safety/performance
324        import numexpr
325        result = numexpr.evaluate(expr, local_dict=local_dict)
326        return str(result.item())
327    except Exception as e:
328        return f"Error calculating expression: {e}"
329    
330# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
331# ─────────────────────────────────────────────── Tool for Date and time ───────────────────────────────────────────────────────────────────
332# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
333
334def get_current_date(_: str = "") -> str:
335    """
336    Returns the current date and time. Ignoring input.
337    """
338    from datetime import datetime
339    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
340
341# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
342# ─────────────────────────────────────────────── Tool for SQL Database ────────────────────────────────────────────────────────────────────
343# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
344
345
346
347def create_sql_agent_function(db_uri: str, top_k: int = 5):
348    """
349    Creates a full-fledged SQL agent function that can answer natural language questions over a SQL database.
350
351    Args:
352        db_uri (str): The SQLAlchemy database URI, e.g. "sqlite:///Chinook.db"
353        top_k (int): Number of rows to limit in results (default 5)
354
355    Returns:
356        agent_executor: LangChain agent that can .run() or .stream()
357    """
358
359    # 1) Initialize the database + LLM + toolkit
360    db = SQLDatabase.from_uri(db_uri)
361    llm = ChatGroq(model="meta-llama/llama-4-maverick-17b-128e-instruct", streaming=False, temperature=0)
362    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
363
364    # 2) Prompt with all required variables declared AND used
365    prompt = PromptTemplate(
366        template="""
367        You are an agent designed to interact with a SQL database.
368        Given the user question below, first generate a syntactically correct {dialect} query.
369        Then look at the results of that query, and return the answer.
370        Always limit to at most {top_k} rows unless the user specifies otherwise.
371        If you encounter an error, rewrite your SQL and retry.
372        DO NOT issue any INSERT/UPDATE/DELETE/DROP/ statements.
373        DO NOT try to create new database tables or columns when user has not asked for.
374        Always inspect the schema before querying.
375
376        Available tools: {tools}
377        Tool names: {tool_names}
378
379        User question: {input}
380
381        {agent_scratchpad}
382        """.strip(),
383                input_variables=["input", "dialect", "top_k", "agent_scratchpad", "tools", "tool_names"],
384            )
385
386    # 3) Create the agent with prompt + toolkit tools
387    agent_executor = create_sql_agent(
388        llm=llm,
389        toolkit=toolkit,
390        prompt=prompt,
391        verbose=False,
392        # pass top_k dynamically
393        extra_prompt_kwargs={"top_k": str(top_k), "dialect": db.dialect},
394    )
395
396    return agent_executor
397
398def execute_sql(query: str) -> str:
399    """
400    Executes a SQL query against the uploaded SQLite DB (GLOBAL_DB_PATH).
401    Returns a string of results or error.
402    """
403    if DB_PATH is None:
404        return "No database uploaded. Please upload a SQLite file first."
405    
406    print("DB_PATH--------->:", DB_PATH)
407    
408    db_uri = f"sqlite:///{DB_PATH}"
409    agent_executor2 = create_sql_agent_function(db_uri, top_k=5)
410  
411    try:
412        result = agent_executor2.run(query)
413    except Exception as e:
414        result = f"Agent / SQL error: {e}"
415    return result
416    
417# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
418# ─────────────────────────────────────────────── Tool for RAG (Document Intelligence) ─────────────────────────────────────────────────────
419# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
420    
421def rag_index_document(DOC_PATH: str) -> str:
422    """
423    Indexes the given document into the RAG vector store.
424    Supports text files or PDFs. Uses recursive text splitting for better chunking.
425    """
426    global vector_store, rag_chain
427    text = ""
428    
429    # Read text from file
430    if DOC_PATH and DOC_PATH.lower().endswith(".pdf"):
431        doc = fitz.open(DOC_PATH)
432        for page in doc:
433            text += page.get_text()
434    else:
435        with open(DOC_PATH, 'r', encoding='utf-8') as f:
436            text = f.read()
437    
438    # Split text using recursive text splitter
439    text_splitter = RecursiveCharacterTextSplitter(
440        chunk_size=500,  # You can adjust this (e.g., 500-1000)
441        chunk_overlap=100  # Overlap for better context between chunks
442    )
443    
444    # Split into chunks
445    texts = text_splitter.split_text(text)
446    
447    # Create Document objects with metadata
448    docs = [Document(page_content=t, metadata={"source": DOC_PATH}) for t in texts]
449    
450    # Initialize or append to FAISS vector store
451    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
452    
453    if vector_store is None:
454        vector_store = FAISS.from_documents(docs, embeddings)
455    else:
456        vector_store.add_documents(docs)
457        
458    retriever = vector_store.as_retriever(
459        search_type="mmr",
460        search_kwargs={
461            "k": 10,
462            "fetch_k": 10,
463            "lambda_mult": 0.25
464        }
465    )
466    
467    # Build or update the RetrievalQA chain
468    rag_chain = RetrievalQA.from_chain_type(
469        llm=llm,
470        chain_type="stuff",
471        retriever=retriever,
472        return_source_documents=False
473    )
474
475
476def rag_answer(query: str) -> str:
477    """
478    Answers a question using the RAG chain (on indexed documents).
479    """
480    global rag_chain
481    if rag_chain is None:
482        return "No documents indexed. Please upload documents via /upload_doc."
483    try:
484        answer = rag_chain.run(query)
485        return answer
486    except Exception as e:
487        return f"RAG error: {e}"
488
489# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
490# ───────────────────────────────────── Tool for Image (understading, captioning & classification) ─────────────────────────────────────────
491# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
492
493# Vision tools and functions  
494# Load image function
495# def _load_image():
496#     try:
497#         if IMG_PATH.startswith("http"):
498#             res = requests.get(IMG_PATH)
499#             res.raise_for_status()
500#             img = Image.open(BytesIO(res.content))
501#         else:
502#             img = Image.open(IMG_PATH)
503#         return img.convert("RGB")
504#     except Exception as e:
505#         raise RuntimeError(f"Failed to load image: {e}")
506
507
508def _load_image(resize_to=(512, 512)):
509    """
510    Load and resize the image from IMG_PATH.
511    If the image is not valid, raise an error.
512    """    
513    try:
514        if IMG_PATH is None:
515            raise ValueError("No image uploaded. Please upload an image first.")
516            #return "No image uploaded. Please upload an image first."
517        with open(IMG_PATH, "rb") as f:
518            img = Image.open(f)
519            img.verify()  # Verify it's an image
520        img = Image.open(IMG_PATH).convert("RGB")  # Reopen after verify and convert
521        img = img.resize(resize_to)  # resize image to reduce token size
522        return img
523    except UnidentifiedImageError:
524        raise ValueError(f"File at {IMG_PATH} is not a valid image.")
525    except Exception as e:
526        raise ValueError(f"Failed to load image at {IMG_PATH}: {str(e)}")
527
528def _encode_image_to_base64():
529    img = _load_image()
530    buffer = BytesIO()
531    img.save(buffer, format="PNG", optimize=True)  # save optimized PNG
532    return base64.b64encode(buffer.getvalue()).decode("utf-8")
533
534def _call_llama_llm(prompt_text: str) -> str:
535    b64 = _encode_image_to_base64()
536    message = HumanMessage(
537        content=[
538            {"type": "text", "text": prompt_text},
539            {
540                "type": "image_url",
541                "image_url": {
542                    "url": f"data:image/png;base64,{b64}"
543                }
544            }
545        ]
546    )
547    response = llm.invoke([message])
548    return response.content.strip()
549
550def vision_query(task_prompt: str) -> str:
551    try:
552        return _call_llama_llm(task_prompt)
553    except Exception as llama_error:
554        print(f"[LLaMA-4V failed] {llama_error}")
555        try:
556            img = _load_image()
557            return pytesseract.image_to_string(img).strip()
558        except Exception as ocr_error:
559            print(f"[OCR fallback failed] {ocr_error}")
560            return "Unable to process the image or image is not uploaded. Please try again with a different input."
561    
562#### Create LangChain Tools ####
563# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
564# ─────────────────────────────────────────────── Assigning tools as list ──────────────────────────────────────────────────────────────────
565# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
566
567tool_list = [
568    Tool(name="browse", func=web_search, description="Search the web and scrape top results. Uses DuckDuckGo (safe mode) for query. Prefers Playwright for loading pages, with requests/BeautifulSoup as fallback. Filters out any explicit content. Returns JSON with titles, URLs, and page text."),
569    Tool(name="calculate", func=calculate, description="Perform math calculations safely."),
570    Tool(name="date", func=get_current_date, description="Fetch the current date and time."),
571    Tool(name="sql", func=execute_sql, description="Execute SQL query on the uploaded database."),
572    Tool(name="rag", func=rag_answer, description="Answer questions using the uploaded documents with retrieval-augmented generation (RAG)."),
573    Tool(
574    name="vision",
575    func=vision_query,
576    description=(
577        "Perform any image-understanding task—e.g. read text, classify objects, "
578        "generate captions, count or locate items, answer questions about the scene, "
579        "detect NSFW content, etc.—powered by LLaMA 4-Vision. "
580        "If the request is OCR-style and LLaMA fails, it falls back to Tesseract OCR."
581    ),
582    ),
583]
584
585# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
586# ─────────────────────────────────────────────── Added Memory to Agent ────────────────────────────────────────────────────────────────────
587# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
588
589# 1) instantiate with your LLM
590memory = AutoSummaryMemory(
591    llm=llm,
592    max_entries=20,     # when chat ≥20 messages, trigger summary
593    reduce_to=5         # keep only last 5 after summarizing
594)
595
596# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
597# ─────────────────────────────────────────────── Initialize Agent ─────────────────────────────────────────────────────────────────────────
598# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
599
600# Initialize the agent with OpenAI and our tools. We use a zero-shot-react-description agent.
601agent_executor = initialize_agent(
602    tools=tool_list,
603    llm=llm,
604    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
605    memory=memory,
606    verbose=True,
607    handle_parsing_errors=True,
608    #max_iterations=10,
609)
610
611# ─── Streaming & Fallback ─────────────────────────────────────────────────────
612# ─── Streaming helper ────────────────────────────────────────────────────────────
613def run_stream(query: str, data_paths: List[str] = None):
614    """
615    Progressive token‐by‐token streaming from the agent.
616
617    Args:
618      query: The user’s natural-language question.
619      data_paths: List of file paths (DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH).
620    """
621    # If no explicit list passed, rebuild from module globals
622    # if not data_paths:
623    #     data_paths = [DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH]
624    data_paths = [p for p in data_paths if p]
625    print(f"Data paths----------------->: {data_paths}")
626    # Re-inject each into the appropriate global (optional—keeps them current)
627    for path in data_paths:
628        ext = os.path.splitext(path)[1].lower()
629        if ext in {".png", ".jpg", ".jpeg", ".gif"}:
630            globals()['IMG_PATH'] = path
631        elif ext in {".pdf", ".txt", ".doc", ".docx"}:
632            globals()['DOC_PATH'] = path
633        elif ext in {".db", ".sqlite"}:
634            globals()['DB_PATH'] = path
635        else:
636            globals()['OTH_PATH'] = path
637
638    # Stream the agent response
639    hist = get_buffer_string(memory.chat_memory.messages)
640    print("Memory now contains:", memory.chat_memory.messages)
641    for chunk in agent_executor.stream({"input": query}):
642        text = chunk.get("text")
643        if text:
644            yield text
645            
646# # ─── Streaming & Fallback ─────────────────────────────────────────────────────
647# def run_stream(query: str, data: str = None):
648#     """
649#     Progressive token‐by‐token streaming from the agent.
650
651#     Args:
652#       query: The user’s natural-language question.
653#       data: Path to a single uploaded file (image, document, or database).
654#             We will inspect its extension and set the appropriate config variable:
655#               .png/.jpg/.jpeg/.gif → IMG_PATH
656#               .pdf/.txt/.doc/.docx   → DOC_PATH
657#               .db/.sqlite            → DB_PATH
658#               others                 → OTH_PATH
659#     """
660#     global DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH
661#     # 1) If data provided, dispatch into the right config variable
662#     if data:
663#         ext = os.path.splitext(data)[1].lower()
664#         if ext in {".png", ".jpg", ".jpeg", ".gif"}:
665#             IMG_PATH = data
666#             print(f"Image path set to: {IMG_PATH}")
667#         elif ext in {".pdf", ".txt", ".doc", ".docx"}:
668#             DOC_PATH = data
669#             print(f"Document path set to: {DOC_PATH}")
670#         elif ext in {".db", ".sqlite"}:
671#             DB_PATH = data
672#             print(f"Database path set to: {DB_PATH}")
673#         else:
674#             OTH_PATH = data
675#             print(f"Other file path set to: {OTH_PATH}")
676
677#     # 2) Stream the agent’s response
678#     for chunk in agent_executor.stream({"input": query}):
679#         text = chunk.get("text")
680#         if text:
681#             yield text
682
683def run_full(query: str) -> str:
684    """
685    Fallback single‐shot answer (for pure-tool or final completeness).
686    """
687    return agent_executor.run(query)
688
689# Expose for Flask
690class AgentInterface:
691    def __init__(self, executor):
692        self.executor = executor
693    def run_stream(self, q):
694        return run_stream(q)
695    def run_full(self, q):
696        return run_full(q)
697
698agent = AgentInterface(agent_executor)
699
700__all__ = [
701    'agent_executor', 'run_stream', 'run_full',
702    'AgentInterface', 'GLOBAL_DB_PATH', 'rag_index_document'
703]
704
705# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
706# ─────────────────────────────────────────────── Refresh Memory Session ───────────────────────────────────────────────────────────────────
707# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
708# Refresh Memory 
709def refresh_memory():
710    memory.clear()  # clear memory at start of each new session
711    memory.chat_memory.clear()  # clear chat history