token-optime 0.1.0__py3-none-any.whl

src/core/cache.py

@@ -0,0 +1,79 @@
+import hashlib, time, sys, os, logging
+from core.client import _chroma_client
+from chromadb.utils import embedding_functions
+from config import settings, PROJECT_ROOT
+
+logger = logging.getLogger("token")
+
+CHROMA_PATH = os.path.join(PROJECT_ROOT, "storage", "chroma_db")
+TTL_BY_TOOL = {
+    "execute": 300,
+    "ask_document": 86400,
+    "search_all_documents": 86400,
+    "list_indexed_documents": 60,
+    "index_document": 0,   
+    "index_documents_folder": 0,
+}
+TTL_DEFAULT = 3600
+
+_cache_collection = None
+
+def _get_cache():
+    global _cache_collection
+    if _cache_collection is None:
+        embedder = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=settings.embedder)
+        _cache_collection = _chroma_client.get_or_create_collection(name="semantic_cache", embedding_function=embedder)
+    return _cache_collection
+
+def check_cache(query: str, tool_name: str = ""):
+    try:
+        results = _get_cache().query(query_texts=[query], n_results=1)
+        if not results["documents"][0]:
+            return None, 0.0
+        distance = results["distances"][0][0]
+        similarity = 1 - distance
+        if similarity >= 0.8:
+            meta = results["metadatas"][0][0]
+            cached_at = meta.get("cached_at", 0)
+            ttl = TTL_BY_TOOL.get(tool_name, TTL_DEFAULT)
+            age = time.time() - cached_at
+            if age > ttl:
+                logger.info(f"[CACHE] EXPIRED for '{tool_name}' | age={int(age)}s ttl={ttl}s sim={similarity:.3f}")
+                return None, similarity
+            logger.info(f"[CACHE] HIT for '{tool_name}' | sim={similarity:.3f} age={int(age)}s")
+            return meta['answer'], similarity
+        return None, similarity
+    except Exception as e:
+        logger.warning(f"[CACHE] check failed: {e}", exc_info=True)
+        return None, 0.0
+
+
+def store_answer(query, answer,tool_name):
+    query_id = hashlib.sha256(query.strip().lower().encode()).hexdigest()
+    try:
+        _get_cache().upsert(
+            ids=[query_id],
+            documents=[query],
+            metadatas=[{'answer': answer, 'cached_at': time.time(), 'tool_name': tool_name}],
+        )
+    except Exception as e:
+        logger.warning(f"[CACHE] store failed for query: {query} - {e}")
+
+
+def cleanup_cache():
+    try:
+        all_entries = _get_cache().get(include=["metadatas"])
+        now = time.time()
+        expired_ids = []
+        for id_, meta in zip(all_entries["ids"], all_entries["metadatas"]):
+            tool_name = meta.get("tool_name", "")
+            ttl = TTL_BY_TOOL.get(tool_name, TTL_DEFAULT)
+            if now - meta.get("cached_at", 0) > ttl:
+                expired_ids.append(id_)
+        if expired_ids:
+            _get_cache().delete(ids=expired_ids)
+            logger.info(f"[CACHE] cleaned up {len(expired_ids)} expired entries")
+        return len(expired_ids)
+    except Exception as e:
+        logger.warning(f"[CACHE] cleanup failed: {e}")
+        return 0

src/core/client.py

@@ -0,0 +1,107 @@
+from groq import Groq
+from config import settings,PROJECT_ROOT
+import chromadb
+import os,json
+import logging
+
+logger = logging.getLogger("token")
+
+_groq_client = Groq(api_key=settings.groq_api_key)
+CHROMA_PATH = os.path.join(PROJECT_ROOT, "storage", "chroma_db")
+_chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
+
+def fill_args_llm(query,schema):
+    props = schema["properties"]
+    required = schema.get("required", [])
+    schema_summary="\n".join([
+        f"- {name} ({info.get('type','string')}) : {info.get('description','no desc')}"
+        for name,info in props.items()
+    ])
+    
+    prompt = f"""You are a tool argument filler. Given a tool schema and a user query, return ONLY a valid JSON object with the correct arguments.
+        Tool parameters: {schema_summary}
+        Required fields: {required}
+        User query: "{query}"
+
+        Rules:
+        - Return ONLY a JSON object, no explanation, no markdown, no backticks
+        - For file paths: NEVER use ~ or relative paths. Always expand to full absolute path.
+        - Home directory is: {os.path.expanduser('~')}
+        - Desktop is: {os.path.expanduser('~/Desktop')}
+        - Documents is: {os.path.expanduser('~/Documents')}
+        - If user says "Desktop", use: {os.path.expanduser('~/Desktop')}
+        - If user says "home", use: {os.path.expanduser('~')}
+        - Example: "file on Desktop" → "{os.path.expanduser('~/Desktop')}/filename.pdf"
+        JSON:
+    """
+        
+    
+    try:
+        raw = _groq_client.chat.completions.create(
+            model=settings.groq_model,
+            messages = [{"role":"user","content":prompt}],
+            temperature=0,
+            max_tokens=settings.max_response_tokens,
+        ) 
+        usage = raw.usage
+        logger.info(f"[GROQ] fill_args | prompt_tokens={usage.prompt_tokens} completion_tokens={usage.completion_tokens}")
+        
+        response = raw.choices[0].message.content.strip()
+        # Strip markdown code fences if model wrapped the JSON
+        if response.startswith("```"):
+            response = response.split("```")[1]
+            if response.startswith("json"):
+                response = response[4:]
+            response = response.strip()
+        args = json.loads(response)
+        
+        result = {k: v for k, v in args.items() if v is not None}
+        result["_groq_usage"] = {
+            "prompt_tokens": usage.prompt_tokens,
+            "completion_tokens": usage.completion_tokens
+        }
+        return result
+    
+    except json.JSONDecodeError as e:
+        logger.warning(f"[GROQ] failed to parse args JSON: {e} — returning error")
+        return {"_groq_error": f"JSONDecodeError: {e}"}
+    except Exception as e:
+        logger.warning(f"[GROQ] fill_args_llm failed: {e}")
+        return {"_groq_error": str(e)}
+
+
+def expand_query(query: str, tool_descriptions: list[dict]) -> tuple[str, int, int]:
+    """query expands for tools to find them """
+    capped = tool_descriptions[:10]
+    tools_text = "\n".join(
+        f"- {t['function']['name']}: {t['function'].get('description', '')}"
+        for t in capped
+    )
+    
+    try:
+        resp = _groq_client.chat.completions.create(
+            model=settings.groq_model,
+            messages=[
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a query rewriter. Given a user query and a list of available tools, "
+                        "rewrite the query using the exact terminology and phrasing that best matches "
+                        "the tool descriptions. Output ONLY the rewritten query, nothing else."
+                    )
+                },
+                {
+                    "role": "user",
+                    "content": f"Tools:\n{tools_text}\n\nUser query: {query}\n\nRewritten query:"
+                }
+            ],
+            max_tokens=60,
+            temperature=0.0
+        )
+        usage = resp.usage
+        logger.info(f"[GROQ] expand_query | prompt_tokens={usage.prompt_tokens} completion_tokens={usage.completion_tokens}")
+        rewritten = resp.choices[0].message.content.strip()
+        return (rewritten if rewritten else query, usage.prompt_tokens, usage.completion_tokens)
+    except Exception as e:
+        logger.warning(f"[EXPAND_QUERY] failed: {e}")
+        return (query, 0, 0)

src/core/db.py

@@ -0,0 +1,339 @@
+import os
+import time
+import uuid
+import logging
+import threading
+import sqlite3
+import requests
+from config import PROJECT_ROOT,settings
+logger = logging.getLogger("token")
+
+# local SQLite — always works
+try:
+
+    DB_PATH = os.path.join(PROJECT_ROOT, "storage", "token_events.db")
+except Exception:
+    DB_PATH = os.path.join(os.path.expanduser("~"), ".token-optime", "storage", "token_events.db")
+
+
+INGEST_URL = settings.token_ingest_url
+
+CONVERSATION_TIMEOUT_MINUTES = 15
+_current_conversation_id = None
+_last_event_time = None
+_local = threading.local()
+_conv_lock = threading.Lock()
+
+
+def _get_or_create_conversation_id():
+    global _current_conversation_id, _last_event_time
+    with _conv_lock:
+        now = time.time()
+        if (
+            _current_conversation_id is None
+            or _last_event_time is None
+            or (now - _last_event_time) > CONVERSATION_TIMEOUT_MINUTES * 60
+        ):
+            _current_conversation_id = str(uuid.uuid4())[:8]
+        _last_event_time = now
+        return _current_conversation_id
+
+
+def get_connection():
+    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def get_thread_connection():
+    if not hasattr(_local, "conn") or _local.conn is None:
+        _local.conn = get_connection()
+        init_db(_local.conn)
+    return _local.conn
+
+
+def close_thread_connection():
+    if hasattr(_local, "conn") and _local.conn is not None:
+        _local.conn.close()
+        _local.conn = None
+
+
+def init_db(conn):
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS events (
+            id  INTEGER PRIMARY KEY AUTOINCREMENT,
+            user_id     TEXT,
+            timestamp   DATETIME DEFAULT CURRENT_TIMESTAMP,
+            conversation_id     TEXT,
+            tool_name   TEXT NOT NULL,
+            query   TEXT,
+            cache_hit   INTEGER DEFAULT 0,
+            cache_similarity    REAL DEFAULT 0.0,
+            tokens_before_trim  INTEGER DEFAULT 0,
+            tokens_after_trim   INTEGER DEFAULT 0,
+            trim_saved  INTEGER DEFAULT 0,
+            schema_tokens_full   INTEGER DEFAULT 0,
+            schema_tokens_selected INTEGER DEFAULT 0,
+            schema_tokens_saved INTEGER DEFAULT 0,
+            groq_prompt_tokens INTEGER DEFAULT 0,
+            groq_completion_tokens  INTEGER DEFAULT 0,
+            doc_id  TEXT,
+            success INTEGER DEFAULT 1
+        )
+    """)
+    for col in ["user_id", "conversation_id"]:
+        try:
+            conn.execute(f"ALTER TABLE events ADD COLUMN {col} TEXT")
+        except Exception:
+            pass
+    conn.commit()
+    logger.info(f"[DB] SQLite ready at {DB_PATH}")
+
+
+def _post_to_cloud(data: dict):
+    """Send event to EC2 API"""
+    try:
+        requests.post(INGEST_URL,json=data,timeout=3,)
+    except Exception as e:
+        logger.debug(f"[DB] Cloud ingest failed (local SQLite has it): {e}")
+
+
+def insert_event(**kwargs):
+    conn = get_thread_connection()
+    kwargs.setdefault("user_id", os.getenv("TOKEN_USER_ID", "unknown"))
+    kwargs.setdefault("conversation_id", _get_or_create_conversation_id())
+
+    fields = [
+        "user_id", "conversation_id",
+        "tool_name", "query", "cache_hit", "cache_similarity",
+        "tokens_before_trim", "tokens_after_trim", "trim_saved",
+        "schema_tokens_full", "schema_tokens_selected", "schema_tokens_saved",
+        "doc_id", "success",
+        "groq_prompt_tokens", "groq_completion_tokens",
+    ]
+
+    data = {f: kwargs.get(f, None) for f in fields}
+
+    # always write to local SQLite
+    placeholders = ", ".join(["?" for _ in fields])
+    columns = ", ".join(fields)
+    conn.execute(
+        f"INSERT INTO events ({columns}) VALUES ({placeholders})",
+        list(data.values()),
+    )
+    conn.commit()
+
+    _post_to_cloud(data)
+        
+        
+def _fetchall(conn, sql, params=()):
+    rows = conn.execute(sql, params).fetchall()
+    return [dict(r) for r in rows]
+
+
+
+def _fetchone(conn, sql, params=()):
+    row = conn.execute(sql, params).fetchone()
+    return dict(row) if row else {}
+
+         
+def get_summary(conn):
+    return _fetchone(conn, """
+        SELECT
+            COUNT(*) as total_calls,
+            COALESCE(SUM(cache_hit), 0) as cache_hits,
+            COALESCE(ROUND(AVG(cache_hit) * 100, 1), 0.0) as hit_rate_pct,
+            COALESCE(SUM(trim_saved), 0) as total_trim_saved,
+            COALESCE(SUM(schema_tokens_saved), 0) as total_schema_saved,
+            COALESCE(SUM(trim_saved + schema_tokens_saved), 0) as total_tokens_saved
+        FROM events
+    """)
+
+
+
+def get_recent_events(conn, limit=50):
+    is_pg = "psycopg2" in type(conn).__module__
+    if is_pg:
+        return _fetchall(conn, "SELECT * FROM events ORDER BY timestamp DESC LIMIT %s", (limit,))
+    else:
+        return _fetchall(conn, "SELECT * FROM events ORDER BY timestamp DESC LIMIT ?", (limit,))
+
+
+def get_tool_stats(conn):
+    return _fetchall(conn,"""
+        SELECT 
+            tool_name,
+            COUNT(*) AS calls,
+            SUM(cache_hit) AS hits,
+            SUM(trim_saved) AS trim_saved,
+            SUM(schema_tokens_saved) AS schema_saved
+        FROM events
+        GROUP BY tool_name
+        ORDER BY calls DESC
+    """)
+    
+
+def get_token_analysis(conn):
+    rows = _fetchone(conn, """
+        SELECT 
+            COUNT(*) AS total_queries,
+            COALESCE(SUM(schema_tokens_saved), 0) AS schema_saved,
+            COALESCE(SUM(schema_tokens_full), 0) AS schema_full,
+            COALESCE(SUM(schema_tokens_selected), 0) AS schema_selected,
+            COALESCE(SUM(trim_saved), 0) AS trim_saved,
+            COALESCE(SUM(tokens_before_trim), 0) AS tokens_before_trim,
+            COALESCE(SUM(tokens_after_trim), 0) AS tokens_after_trim,
+            COALESCE(SUM(cache_hit), 0) AS cache_hits,
+            COUNT(CASE WHEN cache_hit = 0 THEN 1 END) as cache_misses
+        FROM events
+    """)
+    
+    schema_saved = rows.get("schema_saved",0)
+    schema_full = rows.get("schema_full",0)
+    schema_selected = rows.get("schema_selected",0)
+    trim_saved = rows.get("trim_saved",0)
+    tokens_before = rows.get("tokens_before_trim",0)
+    tokens_after = rows.get("tokens_after_trim",0)
+    cache_hits = rows.get("cache_hits",0)
+    cache_misses = rows.get("cache_misses",0)
+    total_queries = rows.get("total_queries",0)
+    
+    total_saved = schema_saved + trim_saved
+    
+    actual_without= schema_full + tokens_before
+    actual_with = schema_selected + tokens_after
+    
+    pct_saved = round((total_saved / actual_without * 100), 1) if actual_without > 0 else 0.0
+
+    return {
+        "total_queries": total_queries,
+        "cache_hits": cache_hits,
+        "cache_misses": cache_misses,
+        "schema_tokens_without_tom": schema_full,
+        "schema_tokens_with_tom": schema_selected,
+        "schema_saved": schema_saved,
+        "tokens_before_trim": tokens_before,
+        "tokens_after_trim": tokens_after,
+        "trim_saved": trim_saved,
+        "total_saved": total_saved,
+        "actual_without_tom": actual_without,
+        "actual_with_tom": actual_with,
+        "pct_saved": pct_saved,
+    }
+    
+def get_event_count(conn) -> int:
+    """Quick row count for health check."""
+    row = _fetchone(conn, "SELECT COUNT(*) AS cnt FROM events")
+    return row.get("cnt", 0)
+
+def get_conversation_stats(conn, limit=20):
+    is_pg = "psycopg2" in type(conn).__module__
+    if is_pg:
+        return _fetchall(conn, """
+            SELECT
+                conversation_id,
+                MIN(timestamp) AS started_at,
+                COUNT(*) AS total_calls,
+                SUM(cache_hit)  AS cache_hits,
+                SUM(trim_saved) AS trim_saved,
+                SUM(schema_tokens_saved) AS schema_saved,
+                SUM(trim_saved + schema_tokens_saved) AS total_saved
+            FROM events
+            WHERE conversation_id IS NOT NULL
+            GROUP BY conversation_id
+            ORDER BY started_at DESC
+            LIMIT %s
+        """, (limit,))
+    else:
+        return _fetchall(conn, """
+            SELECT
+                conversation_id,
+                MIN(timestamp) AS started_at,
+                COUNT(*) AS total_calls,
+                SUM(cache_hit)  AS cache_hits,
+                SUM(trim_saved) AS trim_saved,
+                SUM(schema_tokens_saved) AS schema_saved,
+                SUM(trim_saved + schema_tokens_saved) AS total_saved
+            FROM events
+            WHERE conversation_id IS NOT NULL
+            GROUP BY conversation_id
+            ORDER BY started_at DESC
+            LIMIT ?
+        """, (limit,))
+
+    
+    
+def get_groq_usage(conn):
+    return _fetchone(conn, """
+        SELECT
+            COALESCE(SUM(groq_prompt_tokens), 0) AS total_prompt_tokens,
+            COALESCE(SUM(groq_completion_tokens), 0)  AS total_completion_tokens,
+            COALESCE(SUM(groq_prompt_tokens + groq_completion_tokens), 0) AS total_groq_tokens
+        FROM events
+    """)
+
+
+def get_timeseries(conn, hours: int = 24):
+    is_pg = "psycopg2" in type(conn).__module__
+    if is_pg:
+        return _fetchall(conn, """
+            SELECT
+                to_char(date_trunc('hour', timestamp), 'YYYY-MM-DD HH24:00') AS hour,
+                COUNT(*) AS total_calls,
+                COALESCE(SUM(cache_hit), 0) AS cache_hits,
+                COALESCE(SUM(schema_tokens_saved), 0) AS schema_saved,
+                COALESCE(SUM(trim_saved), 0) AS trim_saved,
+                COALESCE(SUM(groq_prompt_tokens + groq_completion_tokens), 0) AS groq_tokens
+            FROM events
+            WHERE timestamp >= NOW() - INTERVAL '%s hours'
+            GROUP BY date_trunc('hour', timestamp)
+            ORDER BY date_trunc('hour', timestamp) ASC
+        """, (hours,))
+    else:
+        return _fetchall(conn, """
+            SELECT
+                strftime('%Y-%m-%d %H:00', timestamp) AS hour,
+                COUNT(*) AS total_calls,
+                COALESCE(SUM(cache_hit), 0)  AS cache_hits,
+                COALESCE(SUM(schema_tokens_saved), 0) AS schema_saved,
+                COALESCE(SUM(trim_saved), 0) AS trim_saved,
+                COALESCE(SUM(groq_prompt_tokens + groq_completion_tokens), 0) AS groq_tokens
+            FROM events
+            WHERE timestamp >= datetime('now', ? || ' hours')
+            GROUP BY hour
+            ORDER BY hour ASC
+        """, (f"-{hours}",))
+
+    
+def get_period_summary(conn, period: str = "today"):
+    is_pg = "psycopg2" in type(conn).__module__
+    if is_pg:
+        if period == "today":
+            where = "DATE(timestamp) = CURRENT_DATE"
+        elif period == "yesterday":
+            where = "DATE(timestamp) = CURRENT_DATE - INTERVAL '1 day'"
+        elif period == "week":
+            where = "timestamp >= NOW() - INTERVAL '7 days'"
+        else:
+            where = "1=1"
+    else:
+        if period == "today":
+            where = "date(timestamp) = date('now')"
+        elif period == "yesterday":
+            where = "date(timestamp) = date('now', '-1 day')"
+        elif period == "week":
+            where = "timestamp >= datetime('now', '-7 days')"
+        else:
+            where = "1=1"
+
+    return _fetchone(conn, f"""
+        SELECT
+            COUNT(*) AS total_calls,
+            COALESCE(SUM(cache_hit), 0) AS cache_hits,
+            COALESCE(SUM(schema_tokens_saved + trim_saved), 0) AS total_saved,
+            COALESCE(SUM(groq_prompt_tokens), 0)  AS groq_prompt,
+            COALESCE(SUM(groq_completion_tokens), 0) AS groq_completion
+        FROM events
+        WHERE {where}
+    """)

src/core/document_search.py

@@ -0,0 +1,159 @@
+from config import settings, PROJECT_ROOT
+import os, logging,re
+
+logger = logging.getLogger("token")
+CHROMA_PATH = os.path.join(PROJECT_ROOT, "storage", "chroma_db")
+
+_embeddings = None
+_splitter = None
+_vectorstore = None
+_reranker = None
+
+def _get_vectorstore():
+    global _embeddings, _splitter, _vectorstore, _reranker
+    
+    if _vectorstore is None:
+        
+        from langchain_text_splitters import RecursiveCharacterTextSplitter
+        from langchain_huggingface import HuggingFaceEmbeddings
+        from langchain_chroma import Chroma
+        from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+        from langchain_classic.retrievers.document_compressors.cross_encoder_rerank import CrossEncoderReranker
+        
+        
+        _embeddings = HuggingFaceEmbeddings(model_name=settings.embedder)
+        _splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)
+        _vectorstore = Chroma(collection_name="document", embedding_function=_embeddings, persist_directory=CHROMA_PATH)
+        _reranker_model = HuggingFaceCrossEncoder(model_name="cross-encoder/ms-marco-MiniLM-L-6-v2")
+        _reranker = CrossEncoderReranker(model=_reranker_model, top_n=3)
+        
+    return _vectorstore
+
+
+def list_index_doc():
+    """List the documents that are already indexed. this is called before index doc"""
+    try:
+        result = _get_vectorstore().get(include=['metadatas'])
+        seen = {}
+        for meta in result['metadatas']:
+            doc_id = meta.get("doc_id")
+            source = meta.get("source", "unknown")
+            if doc_id and doc_id not in seen:
+                seen[doc_id] = source
+        return seen
+    except Exception as e:
+        logger.error(f"[LIST_DOC] empty : {e}")
+        return {}
+    
+
+def index_doc(file_path, doc_id):
+    """index the doc uploaded by the user"""
+    from langchain_community.document_loaders import PyMuPDFLoader
+    from langchain_core.documents import Document
+    
+    file_path = os.path.expanduser(file_path)
+    if not file_path or not file_path.strip():
+        raise ValueError("file_path cannot be empty")
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+    if os.path.getsize(file_path) == 0:
+        raise ValueError(f"File is empty: {file_path}")
+    if not doc_id or not doc_id.strip():
+        raise ValueError("doc_id cannot be empty")
+    
+    # Validate file extension — only PDFs are supported
+    if not file_path.lower().endswith(".pdf"):
+        raise ValueError(f"Only PDF files are supported. Got: {os.path.splitext(file_path)[1] or '(no extension)'}")
+    
+    
+    doc_id = re.sub(r"[^a-zA-Z0-9_\-]", "_", doc_id.strip())
+    if not doc_id:
+        raise ValueError("doc_id became empty after sanitization — use alphanumeric characters")
+    
+    vs = _get_vectorstore()
+    loader = PyMuPDFLoader(file_path)
+    pages = loader.load()
+    full_chunks = []
+    chunk_index = 0
+    
+    for page in pages:
+        page_number = page.metadata.get("page", 0) + 1
+        chunks = _splitter.split_text(page.page_content)
+        for chunk in chunks:
+            full_chunks.append(
+                Document(page_content=chunk,
+                    metadata={"doc_id": doc_id, "chunk_index": chunk_index,
+                              "page": page_number, "source": file_path}))
+            chunk_index += 1
+    ids = [f"{doc_id}_{i}" for i in range(len(full_chunks))]
+    try:
+        existing = vs.get(where={"doc_id": doc_id})
+        if existing["ids"]:
+            vs.delete(ids=existing["ids"])
+            logger.info(f"[INDEX] removed {len(existing['ids'])} old chunks - {doc_id}")
+    except Exception as e:
+        logger.warning(f"[INDEX] could not clean old chunks for {doc_id}: {e}")
+    vs.add_documents(full_chunks, ids=ids)
+    return len(full_chunks)
+
+
+
+def search_doc(query, doc_id, top_k=3, rerank=10):
+    from langchain_classic.retrievers import ContextualCompressionRetriever
+    
+    vs = _get_vectorstore()
+    search_kwargs = {'k': rerank}
+    if doc_id:
+        search_kwargs['filter'] = {'doc_id': doc_id}
+    base_retriever = vs.as_retriever(search_kwargs=search_kwargs)
+    compress_retriever = ContextualCompressionRetriever(
+        base_compressor=_reranker,
+        base_retriever=base_retriever
+    )
+    results = compress_retriever.invoke(query)
+    return [
+        {"text": doc.page_content, "chunk_index": doc.metadata.get("chunk_index"), "page": doc.metadata.get("page")}
+        for doc in results[:top_k]
+    ]
+    
+    
+
+def search_all_doc(query, top_k=3, rerank=10):
+    from langchain_classic.retrievers import ContextualCompressionRetriever
+    
+    
+    vs = _get_vectorstore()
+    base_retriever = vs.as_retriever(search_kwargs={'k': rerank})
+    compress_retriever = ContextualCompressionRetriever(
+        base_compressor=_reranker,
+        base_retriever=base_retriever
+    )
+    results = compress_retriever.invoke(query)
+    return [
+        {'text': doc.page_content, 'doc_id': doc.metadata.get("doc_id"),
+         'page': doc.metadata.get("page"), 'source': doc.metadata.get("source")}
+        for doc in results[:top_k]
+    ]
+    
+    
+
+def index_folder(folder_path):
+    results = {}
+    if not os.path.isdir(folder_path):
+        raise ValueError(f"Not a directory: {folder_path}")
+    pdf_files = [f for f in os.listdir(folder_path) if f.lower().endswith(".pdf")]
+    if not pdf_files:
+        return {}
+    for filename in pdf_files:
+        file_path = os.path.join(folder_path, filename)
+        doc_id = os.path.splitext(filename)[0]
+        try:
+            chunks = index_doc(file_path, doc_id)
+            results[doc_id] = chunks
+            logger.info(f"Indexed {doc_id} | {chunks} chunks")
+        except Exception as e:
+            logger.error(f"Failed to index {filename}: {e}")
+            results[doc_id] = 0
+    return results
+
+